<commit_before>use std::fmt;
use std::io::{Cursor, Read, Seek, SeekFrom};
use byteorder::{LittleEndian, ReadBytesExt};
use rmpv::decode::read_value;
use protocol::transport::MessageHeader;
use errors::*;

pub trait FromBytes: Sized {
    fn from_bytes(reader: &mut Read) -> Result<Self>;
}

pub trait BlockLength {
    fn block_length() -> u16;
    fn skip_block(buffer: &mut Seek) -> Result<u64> {
        let length = Self::block_length();
        Ok(buffer.seek(SeekFrom::Current(length as i64))?)
    }
}

pub trait Message {
    fn template_id() -> u16;
    fn schema_id() -> u16;
    fn version() -> u16;
}

pub trait ToMessageHeader {
    fn message_header() -> MessageHeader;
}

#[derive(PartialEq, Default)]
pub struct Data(Vec<u8>);

impl ::std::ops::Deref for Data {
    type Target = Vec<u8>;
    fn deref(&self) -> &Vec<u8> {
        &self.0
    }
}

impl From<Data> for Vec<u8> {
    fn from(data: Data) -> Self {
        data.0
    }
}

impl fmt::Debug for Data {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut reader = Cursor::new(&self.0);
        if let Ok(value) = read_value(&mut reader) {
            // assume all message pack data has to be a map
            // to distinguish non message pack data which still
            // can be parsed somehow
            if value.is_map() {
                return write!(f, "{}", value);
            }
        }

        // default debug output
        write!(f, "Data({:?})", self.0)
    }
}

impl FromBytes for u8 {
    fn from_bytes(reader: &mut Read) -> Result<Self> {
        Ok(reader.read_u8()?)
    }
}

impl FromBytes for i8 {
    fn from_bytes(reader: &mut Read) -> Result<Self> {
        Ok(reader.read_i8()?)
    }
}

impl FromBytes for u16 {
    fn from_bytes(reader: &mut Read) -> Result<Self> {
        Ok(reader.read_u16::<LittleEndian>()?)
    }
}

impl FromBytes for i16 {
    fn from_bytes(reader: &mut Read) -> Result<Self> {
        Ok(reader.read_i16::<LittleEndian>()?)
    }
}

impl FromBytes for u32 {
    fn from_bytes(reader: &mut Read) -> Result<Self> {
        Ok(reader.read_u32::<LittleEndian>()?)
    }
}

impl FromBytes for i32 {
    fn from_bytes(reader: &mut Read) -> Result<Self> {
        Ok(reader.read_i32::<LittleEndian>()?)
    }
}

impl FromBytes for u64 {
    fn from_bytes(reader: &mut Read) -> Result<Self> {
        Ok(reader.read_u64::<LittleEndian>()?)
    }
}

impl FromBytes for i64 {
    fn from_bytes(reader: &mut Read) -> Result<Self> {
        Ok(reader.read_i64::<LittleEndian>()?)
    }
}

impl FromBytes for Data {
    fn from_bytes(reader: &mut Read) -> Result<Self> {
        let length = reader.read_u16::<LittleEndian>()?;
        let mut buffer = Vec::with_capacity(length as usize);
        let mut handle = reader.take(length as u64);
        handle.read_to_end(&mut buffer)?;
        Ok(Data(buffer))
    }
}

impl FromBytes for String {
    fn from_bytes(reader: &mut Read) -> Result<Self> {
        let buffer: Data = FromBytes::from_bytes(reader)?;
        Ok(String::from_utf8(buffer.to_vec())?)
    }
}

impl<T: FromBytes> FromBytes for Vec<T> {
    fn from_bytes(reader: &mut Read) -> Result<Self> {
        let _block_length = reader.read_u16::<LittleEndian>()?;
        let num_in_group = reader.read_u8()?;
        let mut group: Vec<T> = Vec::with_capacity(num_in_group as usize);
        for _ in 0..num_in_group {
            group.push(T::from_bytes(reader)?);
        }
        Ok(group)
    }
}

impl FromBytes for bool {
    fn from_bytes(reader: &mut Read) -> Result<Self> {
        match reader.read_u8()? {
            0 => Ok(false),
            1 => Ok(true),
            x => bail!("Invalid boolean encoded as {:?}", x),
        }
    }
}

macro_rules! impl_block_length {
    ($t:ty) => (
        impl BlockLength for $t {
            fn block_length() -> u16 {
                ::std::mem::size_of::<$t>() as u16
            }
        }
    )
}

impl_block_length!(u8);
impl_block_length!(i8);
impl_block_length!(u16);
impl_block_length!(i16);
impl_block_length!(u32);
impl_block_length!(i32);
impl_block_length!(u64);
impl_block_length!(i64);
impl_block_length!(bool);

#[cfg(test)]
#[allow(unused)]
mod tests {

    use super::*;

    #[test]
    fn test_derive_from_bytes_struct() {
        #[derive(FromBytes)]
        struct Foo {
            a: u8,
            b: u16,
            c: u32,
            d: bool,
            e: Data,
            f: String,
        }

        let mut bytes: &[u8] = &[1, 2, 3, 4, 5, 6, 7, 1, 3, 0, 8, 9, 10, 4, 0, 97, 98, 99, 100];
        let foo = Foo::from_bytes(&mut bytes).unwrap();

        assert_eq!(foo.a, 1);
        assert_eq!(foo.b, 770);
        assert_eq!(foo.c, 117835012);
        assert_eq!(foo.d, true);
        assert_eq!(foo.e, Data(vec![8, 9, 10]));
        assert_eq!(foo.f, "abcd");
    }

    #[test]
    fn test_derive_from_bytes_enum() {
        #[derive(Debug, PartialEq, FromBytes)]
        enum Foo {
            A,
            B,
            C = 100,
            D = 122,
            Unknown,
        }

        let mut bytes: &[u8] = &[0, 1, 100, 122, 233];
        let mut next = || Foo::from_bytes(&mut bytes).unwrap();

        assert_eq!(next(), Foo::A);
        assert_eq!(next(), Foo::B);
        assert_eq!(next(), Foo::C);
        assert_eq!(next(), Foo::D);
        assert_eq!(next(), Foo::Unknown);
    }

    #[test]
    fn test_derive_from_bytes_enum_type() {
        #[derive(Debug, PartialEq, FromBytes)]
        #[enum_type = "u32"]
        enum Foo {
            A,
            B = 16909060,
            Unknown,
        }

        let mut bytes: &[u8] = &[0, 0, 0, 0, 4, 3, 2, 1, 1, 0, 0, 0];
        let mut next = || Foo::from_bytes(&mut bytes).unwrap();

        assert_eq!(next(), Foo::A);
        assert_eq!(next(), Foo::B);
        assert_eq!(next(), Foo::Unknown);
    }

    #[test]
    fn test_derive_from_bytes_nested() {
        #[derive(Debug, PartialEq, FromBytes)]
        struct Foo {
            foo: u16,
            bar: Bar,
            baz: Baz,
        }

        #[derive(Debug, PartialEq, FromBytes)]
        enum Bar {
            A,
            B,
            Unknown,
        }

        #[derive(Debug, PartialEq, FromBytes)]
        struct Baz {
            baz: u16,
            bar: Bar,
        }

        let mut bytes: &[u8] = &[12, 0, 1, 0, 2, 0];
        let foo = Foo::from_bytes(&mut bytes).unwrap();

        assert_eq!(foo,
                   Foo {
                       foo: 12,
                       bar: Bar::B,
                       baz: Baz {
                           baz: 512,
                           bar: Bar::A,
                       },
                   });
    }

    #[test]
    fn test_derive_from_bytes_group_encoding() {
        #[derive(Debug, PartialEq, FromBytes)]
        struct Foo {
            foo: u16,
            bar: Vec<Bar>,
            baz: Data,
        }

        #[derive(Debug, PartialEq, FromBytes)]
        struct Bar {
            bar: u16,
            baz: String,
        }

        let mut bytes: &[u8] = &[12, 0, 2, 0, 3, 1, 0, 3, 0, 102, 111, 111, 2, 0, 3, 0, 98, 97, 114, 3, 0, 3, 0, 98, 97, 122, 3, 0, 1, 2,
                                 3];
        let foo = Foo::from_bytes(&mut bytes).unwrap();

        assert_eq!(foo,
                   Foo {
                       foo: 12,
                       bar: vec![Bar {
                                     bar: 1,
                                     baz: "foo".to_string(),
                                 },
                                 Bar {
                                     bar: 2,
                                     baz: "bar".to_string(),
                                 },
                                 Bar {
                                     bar: 3,
                                     baz: "baz".to_string(),
                                 }],
                       baz: Data(vec![1, 2, 3]),
                   });
    }

    #[test]
    fn test_derive_enum_default() {
        #[derive(Debug, PartialEq, EnumDefault)]
        enum Foo {
            A,
            B,
            Unknown,
        }

        assert_eq!(Foo::default(), Foo::Unknown);

        #[derive(Debug, PartialEq, Default)]
        struct Bar {
            a: u16,
            b: Foo,
        }

        assert_eq!(Bar::default(),
                   Bar {
                       a: 0,
                       b: Foo::Unknown,
                   });
    }

    #[test]
    fn test_derive_block_length_struct() {
        #[derive(BlockLength)]
        struct Foo {
            a: u8,
            b: u64,
            c: Vec<Bar>,
            d: Data,
            e: String,
        }

        struct Bar {
            a: u16,
        }

        assert_eq!(Foo::block_length(), 9);
    }

    #[test]
    fn test_derive_block_length_enum() {
        #[derive(BlockLength)]
        enum Foo {
            A,
            B,
            Unknown,
        }

        assert_eq!(Foo::block_length(), 1);
    }

    #[test]
    fn test_derive_block_length_enum_type() {
        #[derive(FromBytes, BlockLength)]
        #[enum_type = "u64"]
        enum Bar {
            A,
            B,
            Unknown,
        }

        assert_eq!(Bar::block_length(), 8);
    }

    #[test]
    fn test_derive_block_length_nested() {
        #[derive(Debug, PartialEq, BlockLength)]
        struct Foo {
            foo: u16,
            bar: Bar,
            baz: Baz,
        }

        #[derive(Debug, PartialEq, BlockLength)]
        #[enum_type = "u32"]
        enum Bar {
            A,
            B,
            Unknown,
        }

        #[derive(Debug, PartialEq, BlockLength)]
        struct Baz {
            baz: u16,
            bar: Bar,
        }

        assert_eq!(Foo::block_length(), 12);

    }

}
<commit_msg>feat(protocol): decode embeded message pack payload<commit_after>use std::fmt;
use std::io::{Cursor, Read, Seek, SeekFrom};
use byteorder::{LittleEndian, ReadBytesExt};
use rmpv::decode::read_value;
use rmpv::{Value, Utf8String};
use protocol::transport::MessageHeader;
use errors::*;

pub trait FromBytes: Sized {
    fn from_bytes(reader: &mut Read) -> Result<Self>;
}

pub trait BlockLength {
    fn block_length() -> u16;
    fn skip_block(buffer: &mut Seek) -> Result<u64> {
        let length = Self::block_length();
        Ok(buffer.seek(SeekFrom::Current(length as i64))?)
    }
}

pub trait Message {
    fn template_id() -> u16;
    fn schema_id() -> u16;
    fn version() -> u16;
}

pub trait ToMessageHeader {
    fn message_header() -> MessageHeader;
}

#[derive(PartialEq, Default)]
pub struct Data(Vec<u8>);

impl ::std::ops::Deref for Data {
    type Target = Vec<u8>;
    fn deref(&self) -> &Vec<u8> {
        &self.0
    }
}

impl From<Data> for Vec<u8> {
    fn from(data: Data) -> Self {
        data.0
    }
}

impl fmt::Debug for Data {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut reader = Cursor::new(&self.0);
        if let Ok(value) = read_value(&mut reader) {
            // assume all message pack data has to be a map
            // to distinguish non message pack data which still
            // can be parsed somehow
            if value.is_map() {
                write!(f, "{}", value)?;

                if let Some(values) = value.as_map() {
                    let payload_key = Value::String(Utf8String::from("payload"));
                    let payload = values.iter().find(|&&(ref key, _)| key == &payload_key);

                    if let Some(&(_, Value::Binary(ref bytes))) = payload {
                        let mut reader = Cursor::new(bytes);
                        if let Ok(value) = read_value(&mut reader) {
                            write!(f, ", payload (decoded): {}", value)?;
                        }
                    }
                }

                return Ok(());
            }
        }

        // default debug output
        write!(f, "Data({:?})", self.0)
    }
}

impl FromBytes for u8 {
    fn from_bytes(reader: &mut Read) -> Result<Self> {
        Ok(reader.read_u8()?)
    }
}

impl FromBytes for i8 {
    fn from_bytes(reader: &mut Read) -> Result<Self> {
        Ok(reader.read_i8()?)
    }
}

impl FromBytes for u16 {
    fn from_bytes(reader: &mut Read) -> Result<Self> {
        Ok(reader.read_u16::<LittleEndian>()?)
    }
}

impl FromBytes for i16 {
    fn from_bytes(reader: &mut Read) -> Result<Self> {
        Ok(reader.read_i16::<LittleEndian>()?)
    }
}

impl FromBytes for u32 {
    fn from_bytes(reader: &mut Read) -> Result<Self> {
        Ok(reader.read_u32::<LittleEndian>()?)
    }
}

impl FromBytes for i32 {
    fn from_bytes(reader: &mut Read) -> Result<Self> {
        Ok(reader.read_i32::<LittleEndian>()?)
    }
}

impl FromBytes for u64 {
    fn from_bytes(reader: &mut Read) -> Result<Self> {
        Ok(reader.read_u64::<LittleEndian>()?)
    }
}

impl FromBytes for i64 {
    fn from_bytes(reader: &mut Read) -> Result<Self> {
        Ok(reader.read_i64::<LittleEndian>()?)
    }
}

impl FromBytes for Data {
    fn from_bytes(reader: &mut Read) -> Result<Self> {
        let length = reader.read_u16::<LittleEndian>()?;
        let mut buffer = Vec::with_capacity(length as usize);
        let mut handle = reader.take(length as u64);
        handle.read_to_end(&mut buffer)?;
        Ok(Data(buffer))
    }
}

impl FromBytes for String {
    fn from_bytes(reader: &mut Read) -> Result<Self> {
        let buffer: Data = FromBytes::from_bytes(reader)?;
        Ok(String::from_utf8(buffer.to_vec())?)
    }
}

impl<T: FromBytes> FromBytes for Vec<T> {
    fn from_bytes(reader: &mut Read) -> Result<Self> {
        let _block_length = reader.read_u16::<LittleEndian>()?;
        let num_in_group = reader.read_u8()?;
        let mut group: Vec<T> = Vec::with_capacity(num_in_group as usize);
        for _ in 0..num_in_group {
            group.push(T::from_bytes(reader)?);
        }
        Ok(group)
    }
}

impl FromBytes for bool {
    fn from_bytes(reader: &mut Read) -> Result<Self> {
        match reader.read_u8()? {
            0 => Ok(false),
            1 => Ok(true),
            x => bail!("Invalid boolean encoded as {:?}", x),
        }
    }
}

macro_rules! impl_block_length {
    ($t:ty) => (
        impl BlockLength for $t {
            fn block_length() -> u16 {
                ::std::mem::size_of::<$t>() as u16
            }
        }
    )
}

impl_block_length!(u8);
impl_block_length!(i8);
impl_block_length!(u16);
impl_block_length!(i16);
impl_block_length!(u32);
impl_block_length!(i32);
impl_block_length!(u64);
impl_block_length!(i64);
impl_block_length!(bool);

#[cfg(test)]
#[allow(unused)]
mod tests {

    use super::*;

    #[test]
    fn test_derive_from_bytes_struct() {
        #[derive(FromBytes)]
        struct Foo {
            a: u8,
            b: u16,
            c: u32,
            d: bool,
            e: Data,
            f: String,
        }

        let mut bytes: &[u8] = &[1, 2, 3, 4, 5, 6, 7, 1, 3, 0, 8, 9, 10, 4, 0, 97, 98, 99, 100];
        let foo = Foo::from_bytes(&mut bytes).unwrap();

        assert_eq!(foo.a, 1);
        assert_eq!(foo.b, 770);
        assert_eq!(foo.c, 117835012);
        assert_eq!(foo.d, true);
        assert_eq!(foo.e, Data(vec![8, 9, 10]));
        assert_eq!(foo.f, "abcd");
    }

    #[test]
    fn test_derive_from_bytes_enum() {
        #[derive(Debug, PartialEq, FromBytes)]
        enum Foo {
            A,
            B,
            C = 100,
            D = 122,
            Unknown,
        }

        let mut bytes: &[u8] = &[0, 1, 100, 122, 233];
        let mut next = || Foo::from_bytes(&mut bytes).unwrap();

        assert_eq!(next(), Foo::A);
        assert_eq!(next(), Foo::B);
        assert_eq!(next(), Foo::C);
        assert_eq!(next(), Foo::D);
        assert_eq!(next(), Foo::Unknown);
    }

    #[test]
    fn test_derive_from_bytes_enum_type() {
        #[derive(Debug, PartialEq, FromBytes)]
        #[enum_type = "u32"]
        enum Foo {
            A,
            B = 16909060,
            Unknown,
        }

        let mut bytes: &[u8] = &[0, 0, 0, 0, 4, 3, 2, 1, 1, 0, 0, 0];
        let mut next = || Foo::from_bytes(&mut bytes).unwrap();

        assert_eq!(next(), Foo::A);
        assert_eq!(next(), Foo::B);
        assert_eq!(next(), Foo::Unknown);
    }

    #[test]
    fn test_derive_from_bytes_nested() {
        #[derive(Debug, PartialEq, FromBytes)]
        struct Foo {
            foo: u16,
            bar: Bar,
            baz: Baz,
        }

        #[derive(Debug, PartialEq, FromBytes)]
        enum Bar {
            A,
            B,
            Unknown,
        }

        #[derive(Debug, PartialEq, FromBytes)]
        struct Baz {
            baz: u16,
            bar: Bar,
        }

        let mut bytes: &[u8] = &[12, 0, 1, 0, 2, 0];
        let foo = Foo::from_bytes(&mut bytes).unwrap();

        assert_eq!(foo,
                   Foo {
                       foo: 12,
                       bar: Bar::B,
                       baz: Baz {
                           baz: 512,
                           bar: Bar::A,
                       },
                   });
    }

    #[test]
    fn test_derive_from_bytes_group_encoding() {
        #[derive(Debug, PartialEq, FromBytes)]
        struct Foo {
            foo: u16,
            bar: Vec<Bar>,
            baz: Data,
        }

        #[derive(Debug, PartialEq, FromBytes)]
        struct Bar {
            bar: u16,
            baz: String,
        }

        let mut bytes: &[u8] = &[12, 0, 2, 0, 3, 1, 0, 3, 0, 102, 111, 111, 2, 0, 3, 0, 98, 97, 114, 3, 0, 3, 0, 98, 97, 122, 3, 0, 1, 2,
                                 3];
        let foo = Foo::from_bytes(&mut bytes).unwrap();

        assert_eq!(foo,
                   Foo {
                       foo: 12,
                       bar: vec![Bar {
                                     bar: 1,
                                     baz: "foo".to_string(),
                                 },
                                 Bar {
                                     bar: 2,
                                     baz: "bar".to_string(),
                                 },
                                 Bar {
                                     bar: 3,
                                     baz: "baz".to_string(),
                                 }],
                       baz: Data(vec![1, 2, 3]),
                   });
    }

    #[test]
    fn test_derive_enum_default() {
        #[derive(Debug, PartialEq, EnumDefault)]
        enum Foo {
            A,
            B,
            Unknown,
        }

        assert_eq!(Foo::default(), Foo::Unknown);

        #[derive(Debug, PartialEq, Default)]
        struct Bar {
            a: u16,
            b: Foo,
        }

        assert_eq!(Bar::default(),
                   Bar {
                       a: 0,
                       b: Foo::Unknown,
                   });
    }

    #[test]
    fn test_derive_block_length_struct() {
        #[derive(BlockLength)]
        struct Foo {
            a: u8,
            b: u64,
            c: Vec<Bar>,
            d: Data,
            e: String,
        }

        struct Bar {
            a: u16,
        }

        assert_eq!(Foo::block_length(), 9);
    }

    #[test]
    fn test_derive_block_length_enum() {
        #[derive(BlockLength)]
        enum Foo {
            A,
            B,
            Unknown,
        }

        assert_eq!(Foo::block_length(), 1);
    }

    #[test]
    fn test_derive_block_length_enum_type() {
        #[derive(FromBytes, BlockLength)]
        #[enum_type = "u64"]
        enum Bar {
            A,
            B,
            Unknown,
        }

        assert_eq!(Bar::block_length(), 8);
    }

    #[test]
    fn test_derive_block_length_nested() {
        #[derive(Debug, PartialEq, BlockLength)]
        struct Foo {
            foo: u16,
            bar: Bar,
            baz: Baz,
        }

        #[derive(Debug, PartialEq, BlockLength)]
        #[enum_type = "u32"]
        enum Bar {
            A,
            B,
            Unknown,
        }

        #[derive(Debug, PartialEq, BlockLength)]
        struct Baz {
            baz: u16,
            bar: Bar,
        }

        assert_eq!(Foo::block_length(), 12);

    }

}
<|endoftext|>


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

use crate::{api, Error, Result};
use async_std::{fs, io};
use slog_scope::{debug, error};
use std::{
    convert::{TryFrom, TryInto},
    path::Path,
};
use surf::{
    http::headers,
    middleware::{self, Middleware},
    Body, Request, Response, StatusCode,
};

struct API;

#[surf::utils::async_trait]
impl Middleware for API {
    async fn handle(
        &self,
        mut req: Request,
        client: surf::Client,
        next: middleware::Next<'_>,
    ) -> surf::Result<Response> {
        req.insert_header(headers::USER_AGENT, "updatehub/next");
        req.insert_header(headers::CONTENT_TYPE, "application/json");
        req.insert_header("api-content-type", "application/vnd.updatehub-v1+json");
        Ok(next.run(req, client).await?)
    }
}

pub struct Client<'a> {
    client: surf::Client,
    server: &'a str,
}

pub async fn get<W>(url: &str, handle: &mut W) -> Result<()>
where
    W: io::Write + Unpin,
{
    validate_url(url)?;
    save_body_to(surf::get(url), handle).await
}

async fn save_body_to<W>(req: surf::RequestBuilder, handle: &mut W) -> Result<()>
where
    W: io::Write + Unpin,
{
    use io::prelude::{ReadExt, WriteExt};
    use std::str::FromStr;

    let mut rep = req.await?;
    if !rep.status().is_success() {
        return Err(Error::InvalidStatusResponse(rep.status()));
    }

    let mut written: f32 = 0.;
    let mut threshold = 10;
    let length = match rep.header(headers::CONTENT_LENGTH) {
        Some(v) => usize::from_str(v.as_str())?,
        None => 0,
    };

    loop {
        let mut buf = [0; 4096];
        let read = rep.read(&mut buf).await?;
        if read == 0 {
            break;
        }
        handle.write_all(&buf[..read]).await?;
        if length > 0 {
            written += read as f32 / (length as f32 / 100.);
            if written as usize >= threshold {
                threshold += 20;
                debug!("{}% of the file has been downloaded", std::cmp::min(written as usize, 100));
            }
        }
    }
    debug!("100% of the file has been downloaded");

    Ok(())
}

impl<'a> Client<'a> {
    pub fn new(server: &'a str) -> Self {
        Self { server, client: surf::Client::new().with(API) }
    }

    pub async fn probe(
        &self,
        num_retries: usize,
        firmware: api::FirmwareMetadata<'_>,
    ) -> Result<api::ProbeResponse> {
        validate_url(self.server)?;

        let mut response = self
            .client
            .post(&format!("{}/upgrades", &self.server))
            .header("api-retries", num_retries.to_string())
            .body(Body::from_json(&firmware)?)
            .await?;

        match response.status() {
            StatusCode::NotFound => Ok(api::ProbeResponse::NoUpdate),
            StatusCode::Ok => {
                match response
                    .header("add-extra-poll")
                    .map(|extra_poll| extra_poll.as_str())
                    .and_then(|extra_poll| extra_poll.parse().ok())
                {
                    Some(extra_poll) => Ok(api::ProbeResponse::ExtraPoll(extra_poll)),
                    None => {
                        let signature =
                            response.header("UH-Signature").map(TryInto::try_into).transpose()?;
                        Ok(api::ProbeResponse::Update(
                            api::UpdatePackage::parse(&response.body_bytes().await?)?,
                            signature,
                        ))
                    }
                }
            }
            s => Err(Error::InvalidStatusResponse(s)),
        }
    }

    pub async fn download_object(
        &self,
        product_uid: &str,
        package_uid: &str,
        download_dir: &Path,
        object: &str,
    ) -> Result<()> {
        validate_url(self.server)?;

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

        if !download_dir.exists() {
            fs::create_dir_all(download_dir).await.map_err(|e| {
                error!("fail to create {:?} directory, error: {}", download_dir, e);
                e
            })?;
        }

        let file = download_dir.join(object);
        if file.exists() {
            request = request
                .header("RANGE", format!("bytes={}-", file.metadata()?.len().saturating_sub(1)));
        }

        let mut file = fs::OpenOptions::new().create(true).append(true).open(&file).await?;

        save_body_to(request, &mut file).await
    }

    pub async fn report(
        &self,
        state: &str,
        firmware: api::FirmwareMetadata<'_>,
        package_uid: &str,
        previous_state: Option<&str>,
        error_message: Option<String>,
        current_log: Option<String>,
    ) -> Result<()> {
        validate_url(self.server)?;

        #[derive(serde::Serialize)]
        #[serde(rename_all = "kebab-case")]
        struct Payload<'a> {
            #[serde(rename = "status")]
            state: &'a str,
            #[serde(flatten)]
            firmware: api::FirmwareMetadata<'a>,
            package_uid: &'a str,
            #[serde(skip_serializing_if = "Option::is_none")]
            previous_state: Option<&'a str>,
            #[serde(skip_serializing_if = "Option::is_none")]
            error_message: Option<String>,
            #[serde(skip_serializing_if = "Option::is_none")]
            current_log: Option<String>,
        }

        let payload =
            Payload { state, firmware, package_uid, previous_state, error_message, current_log };

        self.client
            .post(&format!("{}/report", &self.server))
            .body(Body::from_json(&payload)?)
            .await?;
        Ok(())
    }
}

impl TryFrom<&headers::HeaderValues> for api::Signature {
    type Error = Error;

    fn try_from(value: &headers::HeaderValues) -> Result<Self> {
        let value = value.as_str();

        // Workarround for https://github.com/sfackler/rust-openssl/issues/1325
        if value.is_empty() {
            return Ok(Self::from_base64_str("")?);
        }

        Ok(Self::from_base64_str(value)?)
    }
}

fn validate_url(url: &str) -> surf::Result<()> {
    surf::http::Url::parse(url)?;
    Ok(())
}
<commit_msg>cloud-sdk: Change user agent header as we are close to release<commit_after>// Copyright (C) 2020 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0

use crate::{api, Error, Result};
use async_std::{fs, io};
use slog_scope::{debug, error};
use std::{
    convert::{TryFrom, TryInto},
    path::Path,
};
use surf::{
    http::headers,
    middleware::{self, Middleware},
    Body, Request, Response, StatusCode,
};

struct API;

#[surf::utils::async_trait]
impl Middleware for API {
    async fn handle(
        &self,
        mut req: Request,
        client: surf::Client,
        next: middleware::Next<'_>,
    ) -> surf::Result<Response> {
        req.insert_header(headers::USER_AGENT, "updatehub/2.0 Linux");
        req.insert_header(headers::CONTENT_TYPE, "application/json");
        req.insert_header("api-content-type", "application/vnd.updatehub-v1+json");
        Ok(next.run(req, client).await?)
    }
}

pub struct Client<'a> {
    client: surf::Client,
    server: &'a str,
}

pub async fn get<W>(url: &str, handle: &mut W) -> Result<()>
where
    W: io::Write + Unpin,
{
    validate_url(url)?;
    save_body_to(surf::get(url), handle).await
}

async fn save_body_to<W>(req: surf::RequestBuilder, handle: &mut W) -> Result<()>
where
    W: io::Write + Unpin,
{
    use io::prelude::{ReadExt, WriteExt};
    use std::str::FromStr;

    let mut rep = req.await?;
    if !rep.status().is_success() {
        return Err(Error::InvalidStatusResponse(rep.status()));
    }

    let mut written: f32 = 0.;
    let mut threshold = 10;
    let length = match rep.header(headers::CONTENT_LENGTH) {
        Some(v) => usize::from_str(v.as_str())?,
        None => 0,
    };

    loop {
        let mut buf = [0; 4096];
        let read = rep.read(&mut buf).await?;
        if read == 0 {
            break;
        }
        handle.write_all(&buf[..read]).await?;
        if length > 0 {
            written += read as f32 / (length as f32 / 100.);
            if written as usize >= threshold {
                threshold += 20;
                debug!("{}% of the file has been downloaded", std::cmp::min(written as usize, 100));
            }
        }
    }
    debug!("100% of the file has been downloaded");

    Ok(())
}

impl<'a> Client<'a> {
    pub fn new(server: &'a str) -> Self {
        Self { server, client: surf::Client::new().with(API) }
    }

    pub async fn probe(
        &self,
        num_retries: usize,
        firmware: api::FirmwareMetadata<'_>,
    ) -> Result<api::ProbeResponse> {
        validate_url(self.server)?;

        let mut response = self
            .client
            .post(&format!("{}/upgrades", &self.server))
            .header("api-retries", num_retries.to_string())
            .body(Body::from_json(&firmware)?)
            .await?;

        match response.status() {
            StatusCode::NotFound => Ok(api::ProbeResponse::NoUpdate),
            StatusCode::Ok => {
                match response
                    .header("add-extra-poll")
                    .map(|extra_poll| extra_poll.as_str())
                    .and_then(|extra_poll| extra_poll.parse().ok())
                {
                    Some(extra_poll) => Ok(api::ProbeResponse::ExtraPoll(extra_poll)),
                    None => {
                        let signature =
                            response.header("UH-Signature").map(TryInto::try_into).transpose()?;
                        Ok(api::ProbeResponse::Update(
                            api::UpdatePackage::parse(&response.body_bytes().await?)?,
                            signature,
                        ))
                    }
                }
            }
            s => Err(Error::InvalidStatusResponse(s)),
        }
    }

    pub async fn download_object(
        &self,
        product_uid: &str,
        package_uid: &str,
        download_dir: &Path,
        object: &str,
    ) -> Result<()> {
        validate_url(self.server)?;

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

        if !download_dir.exists() {
            fs::create_dir_all(download_dir).await.map_err(|e| {
                error!("fail to create {:?} directory, error: {}", download_dir, e);
                e
            })?;
        }

        let file = download_dir.join(object);
        if file.exists() {
            request = request
                .header("RANGE", format!("bytes={}-", file.metadata()?.len().saturating_sub(1)));
        }

        let mut file = fs::OpenOptions::new().create(true).append(true).open(&file).await?;

        save_body_to(request, &mut file).await
    }

    pub async fn report(
        &self,
        state: &str,
        firmware: api::FirmwareMetadata<'_>,
        package_uid: &str,
        previous_state: Option<&str>,
        error_message: Option<String>,
        current_log: Option<String>,
    ) -> Result<()> {
        validate_url(self.server)?;

        #[derive(serde::Serialize)]
        #[serde(rename_all = "kebab-case")]
        struct Payload<'a> {
            #[serde(rename = "status")]
            state: &'a str,
            #[serde(flatten)]
            firmware: api::FirmwareMetadata<'a>,
            package_uid: &'a str,
            #[serde(skip_serializing_if = "Option::is_none")]
            previous_state: Option<&'a str>,
            #[serde(skip_serializing_if = "Option::is_none")]
            error_message: Option<String>,
            #[serde(skip_serializing_if = "Option::is_none")]
            current_log: Option<String>,
        }

        let payload =
            Payload { state, firmware, package_uid, previous_state, error_message, current_log };

        self.client
            .post(&format!("{}/report", &self.server))
            .body(Body::from_json(&payload)?)
            .await?;
        Ok(())
    }
}

impl TryFrom<&headers::HeaderValues> for api::Signature {
    type Error = Error;

    fn try_from(value: &headers::HeaderValues) -> Result<Self> {
        let value = value.as_str();

        // Workarround for https://github.com/sfackler/rust-openssl/issues/1325
        if value.is_empty() {
            return Ok(Self::from_base64_str("")?);
        }

        Ok(Self::from_base64_str(value)?)
    }
}

fn validate_url(url: &str) -> surf::Result<()> {
    surf::http::Url::parse(url)?;
    Ok(())
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::convert::Into;

use crate::common::CodecList;
use crate::data::frame::ArcFrame;
use crate::data::packet::Packet;
use crate::data::params::CodecParams;
use crate::data::value::Value;

use crate::error::*;

/// Used to interact with an encoder.
pub trait Encoder: Send {
    /// Returns the extra data added by an encoder to a codec.
    fn get_extradata(&self) -> Option<Vec<u8>>;
    /// Sends to the encoder a frame to be encoded.
    fn send_frame(&mut self, pkt: &ArcFrame) -> Result<()>;
    /// Returns an encoded packet.
    fn receive_packet(&mut self) -> Result<Packet>;
    /// Tells encoder to clear its internal state.
    fn flush(&mut self) -> Result<()>;

    /// Configures the encoder.
    fn configure(&mut self) -> Result<()>;
    /// Sets an encoder option.
    fn set_option<'a>(&mut self, key: &str, val: Value<'a>) -> Result<()>;
    // fn get_option(&mut self, key: &str) -> Option<Value>;
    //
    /// Sets the parameters associated to a determined codec.
    fn set_params(&mut self, params: &CodecParams) -> Result<()>;
    /// Gets the parameters associated to a determined codec.
    fn get_params(&self) -> Result<CodecParams>;
}

/// Auxiliary structure to encapsulate an encoder object and
/// its additional data.
pub struct Context<E: Encoder> {
    enc: E,
    // TODO: Queue up packets/frames
    // TODO: Store here more information
    // TODO: Have a resource pool
    // format: Format
}

impl<E: Encoder> Context<E> {
    // TODO: More constructors
    /// Retrieves a codec descriptor from a codec list through its name,
    /// creates the relative encoder, and encapsulates it into a new `Context`.
    pub fn by_name<T: Descriptor<OutputEncoder = E> + ?Sized>(
        codecs: &Codecs<T>,
        name: &str,
    ) -> Option<Self> {
        if let Some(builder) = codecs.by_name(name) {
            let enc = builder.create();
            Some(Context { enc })
        } else {
            None
        }
    }

    /// Configures the encoder.
    pub fn configure(&mut self) -> Result<()> {
        self.enc.configure()
    }

    /// Sets the parameters associated to a determined codec.
    pub fn set_params(&mut self, params: &CodecParams) -> Result<()> {
        self.enc.set_params(params)
    }

    /// Gets the parameters associated to a determined codec.
    pub fn get_params(&self) -> Result<CodecParams> {
        self.enc.get_params()
    }

    /// Sets an encoder option.
    pub fn set_option<'a, V>(&mut self, key: &str, val: V) -> Result<()>
    where
        V: Into<Value<'a>>,
    {
        // TODO: support more options
        self.enc.set_option(key, val.into())
    }

    /// Returns the extra data added by an encoder to a codec.
    pub fn get_extradata(&mut self) -> Option<Vec<u8>> {
        self.enc.get_extradata()
    }
    /// Sends to the encoder a frame to be encoded.
    pub fn send_frame(&mut self, frame: &ArcFrame) -> Result<()> {
        self.enc.send_frame(frame)
    }
    /// Returns an encoded packet.
    // TODO: Return an Event?
    pub fn receive_packet(&mut self) -> Result<Packet> {
        self.enc.receive_packet()
    }

    /// Tells encoder to clear its internal state.
    pub fn flush(&mut self) -> Result<()> {
        self.enc.flush()
    }
}

/// Codec descriptor.
///
/// Contains information on a codec and its own encoder.
#[derive(Debug)]
pub struct Descr {
    /// The codec name.
    pub codec: &'static str,
    /// The extended codec name.
    pub name: &'static str,
    /// The codec description.
    pub desc: &'static str,
    /// The codec MIME.
    pub mime: &'static str,
    // TODO more fields regarding capabilities
}

/// Used to get the descriptor of a codec and create its own encoder.
pub trait Descriptor {
    type OutputEncoder: Encoder;

    /// Creates a new encoder for the requested codec.
    fn create(&self) -> Self::OutputEncoder;
    /// Returns the codec descriptor.
    fn describe(&self) -> &Descr;
}

/// A list of codec descriptors.
pub struct Codecs<T: 'static + Descriptor + ?Sized> {
    list: HashMap<&'static str, Vec<&'static T>>,
}

impl<T: Descriptor + ?Sized> CodecList for Codecs<T> {
    type D = T;

    fn new() -> Self {
        Codecs {
            list: HashMap::new(),
        }
    }

    // TODO more lookup functions
    fn by_name(&self, name: &str) -> Option<&'static Self::D> {
        self.list.get(name).map(|descs| descs[0])
    }

    fn append(&mut self, desc: &'static Self::D) {
        let codec_name = desc.describe().codec;

        self.list
            .entry(codec_name)
            .or_insert_with(Vec::new)
            .push(desc);
    }
}

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

    mod dummy {
        use super::super::*;
        use crate::data::pixel::Formaton;
        use std::sync::Arc;

        pub struct Enc {
            state: usize,
            w: Option<usize>,
            h: Option<usize>,
            format: Option<Arc<Formaton>>,
        }

        pub struct Des {
            descr: Descr,
        }

        impl Descriptor for Des {
            type OutputEncoder = Enc;

            fn create(&self) -> Self::OutputEncoder {
                Enc {
                    state: 0,
                    w: None,
                    h: None,
                    format: None,
                }
            }

            fn describe(&self) -> &Descr {
                &self.descr
            }
        }

        impl Encoder for Enc {
            fn configure(&mut self) -> Result<()> {
                if self.h.is_some() && self.w.is_some() && self.format.is_some() {
                    Ok(())
                } else {
                    Err(Error::ConfigurationIncomplete)
                }
            }
            fn get_extradata(&self) -> Option<Vec<u8>> {
                Some(vec![self.state as u8; 1])
            }
            fn send_frame(&mut self, _frame: &ArcFrame) -> Result<()> {
                self.state += 1;
                Ok(())
            }
            fn receive_packet(&mut self) -> Result<Packet> {
                let mut p = Packet::with_capacity(1);

                p.data.push(self.state as u8);

                Ok(p)
            }
            fn set_option<'a>(&mut self, key: &str, val: Value<'a>) -> Result<()> {
                match (key, val) {
                    ("w", Value::U64(v)) => self.w = Some(v as usize),
                    ("h", Value::U64(v)) => self.h = Some(v as usize),
                    ("format", Value::Formaton(f)) => self.format = Some(f),
                    _ => return Err(Error::Unsupported(format!("{} key", key))),
                }

                Ok(())
            }

            fn set_params(&mut self, params: &CodecParams) -> Result<()> {
                use crate::data::params::*;

                if let Some(MediaKind::Video(ref info)) = params.kind {
                    self.w = Some(info.width);
                    self.h = Some(info.height);
                    self.format = info.format.clone();
                }
                Ok(())
            }

            fn get_params(&self) -> Result<CodecParams> {
                use crate::data::params::*;

                if self.w.is_none() || self.w.is_none() || self.format.is_none() {
                    return Err(Error::ConfigurationIncomplete);
                }

                Ok(CodecParams {
                    kind: Some(MediaKind::Video(VideoInfo {
                        height: self.w.unwrap(),
                        width: self.h.unwrap(),
                        format: self.format.clone(),
                    })),
                    codec_id: Some("dummy".to_owned()),
                    extradata: self.get_extradata(),
                    bit_rate: 0,
                    convergence_window: 0,
                    delay: 0,
                })
            }

            fn flush(&mut self) -> Result<()> {
                Ok(())
            }
        }

        pub const DUMMY_DESCR: &Des = &Des {
            descr: Descr {
                codec: "dummy",
                name: "dummy",
                desc: "Dummy encoder",
                mime: "x-application/dummy",
            },
        };
    }

    use self::dummy::DUMMY_DESCR;

    #[test]
    fn lookup() {
        let codecs = Codecs::from_list(&[DUMMY_DESCR]);

        let _enc = codecs.by_name("dummy");
    }
}
<commit_msg>codec: Improve function code<commit_after>use std::collections::HashMap;
use std::convert::Into;

use crate::common::CodecList;
use crate::data::frame::ArcFrame;
use crate::data::packet::Packet;
use crate::data::params::CodecParams;
use crate::data::value::Value;

use crate::error::*;

/// Used to interact with an encoder.
pub trait Encoder: Send {
    /// Returns the extra data added by an encoder to a codec.
    fn get_extradata(&self) -> Option<Vec<u8>>;
    /// Sends to the encoder a frame to be encoded.
    fn send_frame(&mut self, pkt: &ArcFrame) -> Result<()>;
    /// Returns an encoded packet.
    fn receive_packet(&mut self) -> Result<Packet>;
    /// Tells encoder to clear its internal state.
    fn flush(&mut self) -> Result<()>;

    /// Configures the encoder.
    fn configure(&mut self) -> Result<()>;
    /// Sets an encoder option.
    fn set_option<'a>(&mut self, key: &str, val: Value<'a>) -> Result<()>;
    // fn get_option(&mut self, key: &str) -> Option<Value>;
    //
    /// Sets the parameters associated to a determined codec.
    fn set_params(&mut self, params: &CodecParams) -> Result<()>;
    /// Gets the parameters associated to a determined codec.
    fn get_params(&self) -> Result<CodecParams>;
}

/// Auxiliary structure to encapsulate an encoder object and
/// its additional data.
pub struct Context<E: Encoder> {
    enc: E,
    // TODO: Queue up packets/frames
    // TODO: Store here more information
    // TODO: Have a resource pool
    // format: Format
}

impl<E: Encoder> Context<E> {
    // TODO: More constructors
    /// Retrieves a codec descriptor from a codec list through its name,
    /// creates the relative encoder, and encapsulates it into a new `Context`.
    pub fn by_name<T: Descriptor<OutputEncoder = E> + ?Sized>(
        codecs: &Codecs<T>,
        name: &str,
    ) -> Option<Self> {
        codecs.by_name(name).map(|builder| Context {
            enc: builder.create(),
        })
    }

    /// Configures the encoder.
    pub fn configure(&mut self) -> Result<()> {
        self.enc.configure()
    }

    /// Sets the parameters associated to a determined codec.
    pub fn set_params(&mut self, params: &CodecParams) -> Result<()> {
        self.enc.set_params(params)
    }

    /// Gets the parameters associated to a determined codec.
    pub fn get_params(&self) -> Result<CodecParams> {
        self.enc.get_params()
    }

    /// Sets an encoder option.
    pub fn set_option<'a, V>(&mut self, key: &str, val: V) -> Result<()>
    where
        V: Into<Value<'a>>,
    {
        // TODO: support more options
        self.enc.set_option(key, val.into())
    }

    /// Returns the extra data added by an encoder to a codec.
    pub fn get_extradata(&mut self) -> Option<Vec<u8>> {
        self.enc.get_extradata()
    }
    /// Sends to the encoder a frame to be encoded.
    pub fn send_frame(&mut self, frame: &ArcFrame) -> Result<()> {
        self.enc.send_frame(frame)
    }
    /// Returns an encoded packet.
    // TODO: Return an Event?
    pub fn receive_packet(&mut self) -> Result<Packet> {
        self.enc.receive_packet()
    }

    /// Tells encoder to clear its internal state.
    pub fn flush(&mut self) -> Result<()> {
        self.enc.flush()
    }
}

/// Codec descriptor.
///
/// Contains information on a codec and its own encoder.
#[derive(Debug)]
pub struct Descr {
    /// The codec name.
    pub codec: &'static str,
    /// The extended codec name.
    pub name: &'static str,
    /// The codec description.
    pub desc: &'static str,
    /// The codec MIME.
    pub mime: &'static str,
    // TODO more fields regarding capabilities
}

/// Used to get the descriptor of a codec and create its own encoder.
pub trait Descriptor {
    type OutputEncoder: Encoder;

    /// Creates a new encoder for the requested codec.
    fn create(&self) -> Self::OutputEncoder;
    /// Returns the codec descriptor.
    fn describe(&self) -> &Descr;
}

/// A list of codec descriptors.
pub struct Codecs<T: 'static + Descriptor + ?Sized> {
    list: HashMap<&'static str, Vec<&'static T>>,
}

impl<T: Descriptor + ?Sized> CodecList for Codecs<T> {
    type D = T;

    fn new() -> Self {
        Codecs {
            list: HashMap::new(),
        }
    }

    // TODO more lookup functions
    fn by_name(&self, name: &str) -> Option<&'static Self::D> {
        self.list.get(name).map(|descs| descs[0])
    }

    fn append(&mut self, desc: &'static Self::D) {
        let codec_name = desc.describe().codec;

        self.list
            .entry(codec_name)
            .or_insert_with(Vec::new)
            .push(desc);
    }
}

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

    mod dummy {
        use super::super::*;
        use crate::data::pixel::Formaton;
        use std::sync::Arc;

        pub struct Enc {
            state: usize,
            w: Option<usize>,
            h: Option<usize>,
            format: Option<Arc<Formaton>>,
        }

        pub struct Des {
            descr: Descr,
        }

        impl Descriptor for Des {
            type OutputEncoder = Enc;

            fn create(&self) -> Self::OutputEncoder {
                Enc {
                    state: 0,
                    w: None,
                    h: None,
                    format: None,
                }
            }

            fn describe(&self) -> &Descr {
                &self.descr
            }
        }

        impl Encoder for Enc {
            fn configure(&mut self) -> Result<()> {
                if self.h.is_some() && self.w.is_some() && self.format.is_some() {
                    Ok(())
                } else {
                    Err(Error::ConfigurationIncomplete)
                }
            }
            fn get_extradata(&self) -> Option<Vec<u8>> {
                Some(vec![self.state as u8; 1])
            }
            fn send_frame(&mut self, _frame: &ArcFrame) -> Result<()> {
                self.state += 1;
                Ok(())
            }
            fn receive_packet(&mut self) -> Result<Packet> {
                let mut p = Packet::with_capacity(1);

                p.data.push(self.state as u8);

                Ok(p)
            }
            fn set_option<'a>(&mut self, key: &str, val: Value<'a>) -> Result<()> {
                match (key, val) {
                    ("w", Value::U64(v)) => self.w = Some(v as usize),
                    ("h", Value::U64(v)) => self.h = Some(v as usize),
                    ("format", Value::Formaton(f)) => self.format = Some(f),
                    _ => return Err(Error::Unsupported(format!("{} key", key))),
                }

                Ok(())
            }

            fn set_params(&mut self, params: &CodecParams) -> Result<()> {
                use crate::data::params::*;

                if let Some(MediaKind::Video(ref info)) = params.kind {
                    self.w = Some(info.width);
                    self.h = Some(info.height);
                    self.format = info.format.clone();
                }
                Ok(())
            }

            fn get_params(&self) -> Result<CodecParams> {
                use crate::data::params::*;

                if self.w.is_none() || self.w.is_none() || self.format.is_none() {
                    return Err(Error::ConfigurationIncomplete);
                }

                Ok(CodecParams {
                    kind: Some(MediaKind::Video(VideoInfo {
                        height: self.w.unwrap(),
                        width: self.h.unwrap(),
                        format: self.format.clone(),
                    })),
                    codec_id: Some("dummy".to_owned()),
                    extradata: self.get_extradata(),
                    bit_rate: 0,
                    convergence_window: 0,
                    delay: 0,
                })
            }

            fn flush(&mut self) -> Result<()> {
                Ok(())
            }
        }

        pub const DUMMY_DESCR: &Des = &Des {
            descr: Descr {
                codec: "dummy",
                name: "dummy",
                desc: "Dummy encoder",
                mime: "x-application/dummy",
            },
        };
    }

    use self::dummy::DUMMY_DESCR;

    #[test]
    fn lookup() {
        let codecs = Codecs::from_list(&[DUMMY_DESCR]);

        let _enc = codecs.by_name("dummy");
    }
}
<|endoftext|>


<commit_before>#[cfg(feature = "serialization")]
use serde::{Deserialize, Serialize};
use std::fmt;
use std::ops::Range;

use crate::ByteIndex;

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[cfg_attr(feature = "serialization", derive(Deserialize, Serialize))]
#[cfg_attr(feature = "memory_usage", derive(heapsize_derive::HeapSizeOf))]
pub struct Span {
    start: ByteIndex,
    end: ByteIndex,
}

impl Span {
    /// Create a new span from a starting and ending span.
    pub fn new(start: impl Into<ByteIndex>, end: impl Into<ByteIndex>) -> Span {
        let start = start.into();
        let end = end.into();

        assert!(end >= start);

        Span { start, end }
    }

    /// Gives an empty span at the start of a source.
    pub fn initial() -> Span {
        Span::new(0, 0)
    }

    /// Measure the span of a string.
    ///
    /// ```rust
    /// use codespan::{ByteIndex, Span};
    ///
    /// let span = Span::from_str("hello");
    ///
    /// assert_eq!(span, Span::new(0, 5));
    /// ```
    pub fn from_str(s: &str) -> Span {
        Span::new(0, s.len() as u32)
    }

    /// Combine two spans by taking the start of the first span and the end of
    /// the other span.
    ///
    /// ```rust
    /// use codespan::{ByteIndex, Span};
    ///
    /// let span1 = Span::new(0, 4);
    /// let span2 = Span::new(10, 16);
    ///
    /// assert_eq!(Span::merge(span1, span2), Span::new(0, 16));
    /// ```
    pub fn merge(self, other: Span) -> Span {
        Span::new(self.start(), other.end())
    }

    /// Get the starting byte index.
    ///
    /// ```rust
    /// use codespan::{ByteIndex, Span};
    ///
    /// let span = Span::new(0, 4);
    ///
    /// assert_eq!(span.start(), ByteIndex::from(0));
    /// ```
    pub fn start(&self) -> ByteIndex {
        self.start
    }

    /// Get the ending byte index.
    ///
    /// ```rust
    /// use codespan::{ByteIndex, Span};
    ///
    /// let span = Span::new(0, 4);
    ///
    /// assert_eq!(span.end(), ByteIndex::from(4));
    /// ```
    pub fn end(&self) -> ByteIndex {
        self.end
    }
}

impl Default for Span {
    fn default() -> Span {
        Span::initial()
    }
}

impl fmt::Display for Span {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "[{start}, {end})",
            start = self.start(),
            end = self.end(),
        )
    }
}

impl<I> From<Range<I>> for Span
where
    I: Into<ByteIndex>,
{
    fn from(range: Range<I>) -> Span {
        Span::new(range.start, range.end)
    }
}
<commit_msg>Make Span::initial() a const fn<commit_after>#[cfg(feature = "serialization")]
use serde::{Deserialize, Serialize};
use std::fmt;
use std::ops::Range;

use crate::ByteIndex;

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[cfg_attr(feature = "serialization", derive(Deserialize, Serialize))]
#[cfg_attr(feature = "memory_usage", derive(heapsize_derive::HeapSizeOf))]
pub struct Span {
    start: ByteIndex,
    end: ByteIndex,
}

impl Span {
    /// Create a new span from a starting and ending span.
    pub fn new(start: impl Into<ByteIndex>, end: impl Into<ByteIndex>) -> Span {
        let start = start.into();
        let end = end.into();

        assert!(end >= start);

        Span { start, end }
    }

    /// Gives an empty span at the start of a source.
    pub const fn initial() -> Span {
        Span {
            start: ByteIndex(0),
            end: ByteIndex(0),
        }
    }

    /// Measure the span of a string.
    ///
    /// ```rust
    /// use codespan::{ByteIndex, Span};
    ///
    /// let span = Span::from_str("hello");
    ///
    /// assert_eq!(span, Span::new(0, 5));
    /// ```
    pub fn from_str(s: &str) -> Span {
        Span::new(0, s.len() as u32)
    }

    /// Combine two spans by taking the start of the first span and the end of
    /// the other span.
    ///
    /// ```rust
    /// use codespan::{ByteIndex, Span};
    ///
    /// let span1 = Span::new(0, 4);
    /// let span2 = Span::new(10, 16);
    ///
    /// assert_eq!(Span::merge(span1, span2), Span::new(0, 16));
    /// ```
    pub fn merge(self, other: Span) -> Span {
        Span::new(self.start(), other.end())
    }

    /// Get the starting byte index.
    ///
    /// ```rust
    /// use codespan::{ByteIndex, Span};
    ///
    /// let span = Span::new(0, 4);
    ///
    /// assert_eq!(span.start(), ByteIndex::from(0));
    /// ```
    pub fn start(&self) -> ByteIndex {
        self.start
    }

    /// Get the ending byte index.
    ///
    /// ```rust
    /// use codespan::{ByteIndex, Span};
    ///
    /// let span = Span::new(0, 4);
    ///
    /// assert_eq!(span.end(), ByteIndex::from(4));
    /// ```
    pub fn end(&self) -> ByteIndex {
        self.end
    }
}

impl Default for Span {
    fn default() -> Span {
        Span::initial()
    }
}

impl fmt::Display for Span {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "[{start}, {end})",
            start = self.start(),
            end = self.end(),
        )
    }
}

impl<I> From<Range<I>> for Span
where
    I: Into<ByteIndex>,
{
    fn from(range: Range<I>) -> Span {
        Span::new(range.start, range.end)
    }
}
<|endoftext|>


<commit_before>extern crate getopts;
extern crate notify;

use self::notify::{Watcher, RecursiveMode, DebouncedEvent};
use std::sync::mpsc::channel;
use std::time::Duration;
use std::env;
use std::fs;
use std::path;
use std::thread;
use std::process::{Command,Child};

/// watches the debug build, restarts server on build
pub fn watcher (matches: &getopts::Matches) {
    if !matches.opt_present("w") { return }

    let stratis_project = env::var("STRATIS").expect("no \'STRATIS\' path found in environment vars");
    let server_path;
    let server_dest;

    let ffi_path;
    let ffi_dest = env::var("STRATIS_LINK");

    if matches.opt_present("d") {
        server_path = stratis_project.clone() + "\\server\\target\\debug\\";
        server_dest = server_path.clone()+"PAN_stratis.exe";

        ffi_path = stratis_project.clone()+"\\ffi\\target\\debug\\";
    }
    else {
        server_path = stratis_project.clone() + "\\server\\target\\release\\";
        server_dest = server_path.clone()+"\\PAN_stratis.exe";

        ffi_path = stratis_project.clone()+"\\ffi\\target\\release\\";
    }

    let (tx, rx) = channel();
    let mut w = self::notify::watcher(tx,Duration::from_secs(1)).expect("unable to create filesys watcher");

    let _ = w.watch(server_path.clone()+"stratis.exe",RecursiveMode::NonRecursive);

    if let Ok(ref dest) = ffi_dest {
        println!("watching stratis_ffi builds for:{:?}",dest);
        let _ = w.watch(ffi_path.clone()+"stratis_ffi.dll",RecursiveMode::NonRecursive);
    }

    let mut spawn_handle: Option<Child> = None;
    if matches.opt_present("r") {

        if Command::new("cargo")
            .current_dir(stratis_project.clone() + "\\server\\")
            .arg("build")
            .status().expect("failed to build stratis server").success() {
                spawn_handle = spawn(&server_dest);
            }
    }

    for n in rx.iter() {
        let mut new_spawn_handle = None;
        println!("event: {:?}",n);

        match n {
            DebouncedEvent::Write(path) => {
                if path == path::PathBuf::from(server_path.clone()+"stratis.exe") {
                    if let Some(ref mut h) = spawn_handle {
                        if let Ok(_) = h.kill() {
                            println!("process-killed\n");
                            thread::sleep(Duration::from_secs(1));

                            if let Ok(r) = fs::copy(&path,&server_dest) {
                                println!("watcher-copy:{:?}",r);
                                new_spawn_handle = spawn(&server_dest);
                            }
                        }
                    }
                    else { //no child process alive?
                        if let Ok(r) = fs::copy(&path,&server_dest) {
                            println!("watcher-copy:{:?}",r);
                            new_spawn_handle = spawn(&server_dest);
                        }
                    }

                    spawn_handle = new_spawn_handle;
                    if let Some(ref h) = spawn_handle {
                        println!("process-spawned:{:?}\n",h.id());
                    }
                }
                else if path == path::PathBuf::from(ffi_path.clone()+"stratis_ffi.dll") {
                    if let Ok(_) = ffi_dest { //TODO: copy dll to STRATIS_LINK path?
                        println!("copying ffi dll");
                        let r = fs::copy(&path,stratis_project.clone()+"\\unity_ffi\\stratis_ffi.dll"); // we do this because VS is stupid
                        println!("watcher-copy:{:?}",r);

                        let r = fs::copy(&path,stratis_project.clone()+"\\unity_ffi\\Assets\\Plugins\\stratis_ffi.dll");
                        println!("watcher-copy:{:?}",r);
                    }
                }
            },
            _ => {  },
        }
    }

}

fn spawn(path: &str) -> Option<Child> {
    let r = Command::new(path).spawn();
    match r {
        Ok(h) => {
            Some(h)
        },
        Err(e) => {
            println!("spawn-err:{:?}",e);
            None
        }
    }
}
<commit_msg>handle when first being built<commit_after>extern crate getopts;
extern crate notify;

use self::notify::{Watcher, RecursiveMode, DebouncedEvent};
use std::sync::mpsc::channel;
use std::time::Duration;
use std::env;
use std::fs;
use std::path;
use std::thread;
use std::process::{Command,Child};

/// watches the debug build, restarts server on build
pub fn watcher (matches: &getopts::Matches) {
    if !matches.opt_present("w") { return }

    let stratis_project = env::var("STRATIS").expect("no \'STRATIS\' path found in environment vars");
    let server_path;
    let server_dest;

    let ffi_path;
    let ffi_dest = env::var("STRATIS_LINK");

    if matches.opt_present("d") {
        server_path = stratis_project.clone() + "\\server\\target\\debug\\";
        server_dest = server_path.clone()+"PAN_stratis.exe";

        ffi_path = stratis_project.clone()+"\\ffi\\target\\debug\\";
    }
    else {
        server_path = stratis_project.clone() + "\\server\\target\\release\\";
        server_dest = server_path.clone()+"\\PAN_stratis.exe";

        ffi_path = stratis_project.clone()+"\\ffi\\target\\release\\";
    }

    let (tx, rx) = channel();
    let mut w = self::notify::watcher(tx,Duration::from_secs(1)).expect("unable to create filesys watcher");

    let _ = w.watch(server_path.clone()+"stratis.exe",RecursiveMode::NonRecursive);

    if let Ok(ref dest) = ffi_dest {
        println!("watching stratis_ffi builds for:{:?}",dest);
        let _ = w.watch(ffi_path.clone()+"stratis_ffi.dll",RecursiveMode::NonRecursive);
    }

    let mut spawn_handle: Option<Child> = None;
    if matches.opt_present("r") {

        if Command::new("cargo")
            .current_dir(stratis_project.clone() + "\\server\\")
            .arg("build")
            .status().expect("failed to build stratis server").success() {
                spawn_handle = spawn(&server_dest);
            }
    }

    if spawn_handle.is_none() {
        let path = path::PathBuf::from(server_path.clone()+"stratis.exe");
        let r = fs::copy(&path,&server_dest);
        spawn_handle = spawn(&server_dest);

        println!("Missing build, re-copying {:?} & respawning {:?}",r, spawn_handle.is_some());
    }

    for n in rx.iter() {
        let mut new_spawn_handle = None;
        println!("event: {:?}",n);

        match n {
            DebouncedEvent::Write(path) => {
                if path == path::PathBuf::from(server_path.clone()+"stratis.exe") {
                    if let Some(ref mut h) = spawn_handle {
                        if let Ok(_) = h.kill() {
                            println!("process-killed\n");
                            thread::sleep(Duration::from_secs(1));

                            if let Ok(r) = fs::copy(&path,&server_dest) {
                                println!("watcher-copy:{:?}",r);
                                new_spawn_handle = spawn(&server_dest);
                            }
                        }
                    }
                    else { //no child process alive?
                        if let Ok(r) = fs::copy(&path,&server_dest) {
                            println!("watcher-copy:{:?}",r);
                            new_spawn_handle = spawn(&server_dest);
                        }
                    }

                    spawn_handle = new_spawn_handle;
                    if let Some(ref h) = spawn_handle {
                        println!("process-spawned:{:?}\n",h.id());
                    }
                }
                else if path == path::PathBuf::from(ffi_path.clone()+"stratis_ffi.dll") {
                    if let Ok(_) = ffi_dest { //TODO: copy dll to STRATIS_LINK path?
                        println!("copying ffi dll");
                        let r = fs::copy(&path,stratis_project.clone()+"\\unity_ffi\\stratis_ffi.dll"); // we do this because VS is stupid
                        println!("watcher-copy:{:?}",r);

                        let r = fs::copy(&path,stratis_project.clone()+"\\unity_ffi\\Assets\\Plugins\\stratis_ffi.dll");
                        println!("watcher-copy:{:?}",r);
                    }
                }
            },
            _ => {  },
        }
    }

}

fn spawn(path: &str) -> Option<Child> {
    let r = Command::new(path).spawn();
    match r {
        Ok(h) => {
            Some(h)
        },
        Err(e) => {
            println!("spawn-err:{:?}",e);
            None
        }
    }
}
<|endoftext|>


<commit_before>use reqwest::StatusCode;
use std::error;
use std::fmt;

#[derive(Fail, Debug)]
pub enum WalmartError {
  #[fail(display = "{}", _0)]
  Msg(String),

  #[fail(display = "api response error: {}", _0)]
  Api(ApiResponseError),

  #[fail(display = "http error: {}", _0)]
  Reqwest(::reqwest::Error),

  #[fail(display = "url error: {}", _0)]
  UrlError(::reqwest::UrlError),

  #[fail(display = "base64 error: {}", _0)]
  Base64(::base64::DecodeError),

  #[fail(display = "openssl error: {}", _0)]
  OpenSSL(::openssl::error::ErrorStack),

  #[fail(display = "url query serialize error: {}", _0)]
  UrlEncoded(::serde_urlencoded::ser::Error),

  #[fail(display = "io error: {}", _0)]
  Io(::std::io::Error),

  #[fail(display = "zip error: {}", _0)]
  Zip(::zip::result::ZipError),

  #[fail(display = "xml parse error: {}", _0)]
  XmlParse(::xmltree::ParseError),

  #[fail(display = "csv error: {}", _0)]
  Csv(::csv::Error),
}

macro_rules! impl_from {
  ($v:ident($t:ty)) => {
    impl From<$t> for WalmartError {
      fn from(e: $t) -> Self {
        WalmartError::$v(e)
      }
    }
  };
}

impl_from!(Msg(String));
impl_from!(Api(ApiResponseError));
impl_from!(Reqwest(::reqwest::Error));
impl_from!(UrlError(::reqwest::UrlError));
impl_from!(Base64(::base64::DecodeError));
impl_from!(OpenSSL(::openssl::error::ErrorStack));
impl_from!(UrlEncoded(::serde_urlencoded::ser::Error));
impl_from!(Io(::std::io::Error));
impl_from!(Zip(::zip::result::ZipError));
impl_from!(XmlParse(::xmltree::ParseError));
impl_from!(Csv(::csv::Error));

#[derive(Debug)]
pub struct ApiResponseError {
  pub message: String,
  pub status: StatusCode,
  pub body: String,
}

impl fmt::Display for ApiResponseError {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    write!(
      f,
      "API Error: status = '{}', message = '{}'",
      self.status, self.message
    )
  }
}

impl error::Error for ApiResponseError {
  fn description(&self) -> &str {
    "API response error"
  }

  fn cause(&self) -> Option<&error::Error> {
    None
  }
}

pub type WalmartResult<T> = Result<T, WalmartError>;
<commit_msg>WalmartError::should_try_again<commit_after>use reqwest::StatusCode;
use std::error;
use std::fmt;

#[derive(Fail, Debug)]
pub enum WalmartError {
  #[fail(display = "{}", _0)]
  Msg(String),

  #[fail(display = "api response error: {}", _0)]
  Api(ApiResponseError),

  #[fail(display = "http error: {}", _0)]
  Reqwest(::reqwest::Error),

  #[fail(display = "url error: {}", _0)]
  UrlError(::reqwest::UrlError),

  #[fail(display = "base64 error: {}", _0)]
  Base64(::base64::DecodeError),

  #[fail(display = "openssl error: {}", _0)]
  OpenSSL(::openssl::error::ErrorStack),

  #[fail(display = "url query serialize error: {}", _0)]
  UrlEncoded(::serde_urlencoded::ser::Error),

  #[fail(display = "io error: {}", _0)]
  Io(::std::io::Error),

  #[fail(display = "zip error: {}", _0)]
  Zip(::zip::result::ZipError),

  #[fail(display = "xml parse error: {}", _0)]
  XmlParse(::xmltree::ParseError),

  #[fail(display = "csv error: {}", _0)]
  Csv(::csv::Error),
}

impl WalmartError {
  pub fn should_try_again(&self) -> bool {
    match *self {
      WalmartError::Reqwest(ref err) => {
        if let Some(status) = err.status() {
          let code = status.as_u16();
          // 429 Too Many Requests
          code == 429 || code == 500 || code == 503
        } else {
          false
        }
      }
      _ => false,
    }
  }
}

macro_rules! impl_from {
  ($v:ident($t:ty)) => {
    impl From<$t> for WalmartError {
      fn from(e: $t) -> Self {
        WalmartError::$v(e)
      }
    }
  };
}

impl_from!(Msg(String));
impl_from!(Api(ApiResponseError));
impl_from!(Reqwest(::reqwest::Error));
impl_from!(UrlError(::reqwest::UrlError));
impl_from!(Base64(::base64::DecodeError));
impl_from!(OpenSSL(::openssl::error::ErrorStack));
impl_from!(UrlEncoded(::serde_urlencoded::ser::Error));
impl_from!(Io(::std::io::Error));
impl_from!(Zip(::zip::result::ZipError));
impl_from!(XmlParse(::xmltree::ParseError));
impl_from!(Csv(::csv::Error));

#[derive(Debug)]
pub struct ApiResponseError {
  pub message: String,
  pub status: StatusCode,
  pub body: String,
}

impl fmt::Display for ApiResponseError {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    write!(
      f,
      "API Error: status = '{}', message = '{}'",
      self.status, self.message
    )
  }
}

impl error::Error for ApiResponseError {
  fn description(&self) -> &str {
    "API response error"
  }

  fn cause(&self) -> Option<&error::Error> {
    None
  }
}

pub type WalmartResult<T> = Result<T, WalmartError>;
<|endoftext|>


<commit_before>use std::fmt;
use std::marker::PhantomData;
use std::time::Duration;

use actix_codec::{AsyncRead, AsyncWrite};
use actix_connect::{default_connector, Connect as TcpConnect, TcpConnection};
use actix_service::{apply_fn, Service, ServiceExt};
use actix_utils::timeout::{TimeoutError, TimeoutService};
use http::Uri;
use tokio_tcp::TcpStream;

use super::connection::Connection;
use super::error::ConnectError;
use super::pool::{ConnectionPool, Protocol};

#[cfg(feature = "ssl")]
use openssl::ssl::SslConnector;

#[cfg(not(feature = "ssl"))]
type SslConnector = ();

/// Http client connector builde instance.
/// `Connector` type uses builder-like pattern for connector service construction.
pub struct Connector<T, U> {
    connector: T,
    timeout: Duration,
    conn_lifetime: Duration,
    conn_keep_alive: Duration,
    disconnect_timeout: Duration,
    limit: usize,
    #[allow(dead_code)]
    ssl: SslConnector,
    _t: PhantomData<U>,
}

impl Connector<(), ()> {
    pub fn new() -> Connector<
        impl Service<
                Request = TcpConnect<Uri>,
                Response = TcpConnection<Uri, TcpStream>,
                Error = actix_connect::ConnectError,
            > + Clone,
        TcpStream,
    > {
        let ssl = {
            #[cfg(feature = "ssl")]
            {
                use log::error;
                use openssl::ssl::{SslConnector, SslMethod};

                let mut ssl = SslConnector::builder(SslMethod::tls()).unwrap();
                let _ = ssl
                    .set_alpn_protos(b"\x02h2\x08http/1.1")
                    .map_err(|e| error!("Can not set alpn protocol: {:?}", e));
                ssl.build()
            }
            #[cfg(not(feature = "ssl"))]
            {}
        };

        Connector {
            ssl,
            connector: default_connector(),
            timeout: Duration::from_secs(1),
            conn_lifetime: Duration::from_secs(75),
            conn_keep_alive: Duration::from_secs(15),
            disconnect_timeout: Duration::from_millis(3000),
            limit: 100,
            _t: PhantomData,
        }
    }
}

impl<T, U> Connector<T, U> {
    /// Use custom connector.
    pub fn connector<T1, U1>(self, connector: T1) -> Connector<T1, U1>
    where
        U1: AsyncRead + AsyncWrite + fmt::Debug,
        T1: Service<
                Request = TcpConnect<Uri>,
                Response = TcpConnection<Uri, U1>,
                Error = actix_connect::ConnectError,
            > + Clone,
    {
        Connector {
            connector,
            timeout: self.timeout,
            conn_lifetime: self.conn_lifetime,
            conn_keep_alive: self.conn_keep_alive,
            disconnect_timeout: self.disconnect_timeout,
            limit: self.limit,
            ssl: self.ssl,
            _t: PhantomData,
        }
    }
}

impl<T, U> Connector<T, U>
where
    U: AsyncRead + AsyncWrite + fmt::Debug + 'static,
    T: Service<
            Request = TcpConnect<Uri>,
            Response = TcpConnection<Uri, U>,
            Error = actix_connect::ConnectError,
        > + Clone,
{
    /// Connection timeout, i.e. max time to connect to remote host including dns name resolution.
    /// Set to 1 second by default.
    pub fn timeout(mut self, timeout: Duration) -> Self {
        self.timeout = timeout;
        self
    }

    #[cfg(feature = "ssl")]
    /// Use custom `SslConnector` instance.
    pub fn ssl(mut self, connector: SslConnector) -> Self {
        self.ssl = connector;
        self
    }

    /// Set total number of simultaneous connections per type of scheme.
    ///
    /// If limit is 0, the connector has no limit.
    /// The default limit size is 100.
    pub fn limit(mut self, limit: usize) -> Self {
        self.limit = limit;
        self
    }

    /// Set keep-alive period for opened connection.
    ///
    /// Keep-alive period is the period between connection usage. If
    /// the delay between repeated usages of the same connection
    /// exceeds this period, the connection is closed.
    /// Default keep-alive period is 15 seconds.
    pub fn conn_keep_alive(mut self, dur: Duration) -> Self {
        self.conn_keep_alive = dur;
        self
    }

    /// Set max lifetime period for connection.
    ///
    /// Connection lifetime is max lifetime of any opened connection
    /// until it is closed regardless of keep-alive period.
    /// Default lifetime period is 75 seconds.
    pub fn conn_lifetime(mut self, dur: Duration) -> Self {
        self.conn_lifetime = dur;
        self
    }

    /// Set server connection disconnect timeout in milliseconds.
    ///
    /// Defines a timeout for disconnect connection. If a disconnect procedure does not complete
    /// within this time, the socket get dropped. This timeout affects only secure connections.
    ///
    /// To disable timeout set value to 0.
    ///
    /// By default disconnect timeout is set to 3000 milliseconds.
    pub fn disconnect_timeout(mut self, dur: Duration) -> Self {
        self.disconnect_timeout = dur;
        self
    }

    /// Finish configuration process and create connector service.
    pub fn service(
        self,
    ) -> impl Service<Request = Uri, Response = impl Connection, Error = ConnectError> + Clone
    {
        #[cfg(not(feature = "ssl"))]
        {
            let connector = TimeoutService::new(
                self.timeout,
                apply_fn(self.connector, |msg: Uri, srv| srv.call(msg.into()))
                    .map_err(ConnectError::from)
                    .map(|stream| (stream.into_parts().0, Protocol::Http1)),
            )
            .map_err(|e| match e {
                TimeoutError::Service(e) => e,
                TimeoutError::Timeout => ConnectError::Timeout,
            });

            connect_impl::InnerConnector {
                tcp_pool: ConnectionPool::new(
                    connector,
                    self.conn_lifetime,
                    self.conn_keep_alive,
                    None,
                    self.limit,
                ),
            }
        }
        #[cfg(feature = "ssl")]
        {
            const H2: &[u8] = b"h2";
            use actix_connect::ssl::OpensslConnector;

            let ssl_service = TimeoutService::new(
                self.timeout,
                apply_fn(self.connector.clone(), |msg: Uri, srv| srv.call(msg.into()))
                    .map_err(ConnectError::from)
                    .and_then(
                        OpensslConnector::service(self.ssl)
                            .map_err(ConnectError::from)
                            .map(|stream| {
                                let sock = stream.into_parts().0;
                                let h2 = sock
                                    .get_ref()
                                    .ssl()
                                    .selected_alpn_protocol()
                                    .map(|protos| protos.windows(2).any(|w| w == H2))
                                    .unwrap_or(false);
                                if h2 {
                                    (sock, Protocol::Http2)
                                } else {
                                    (sock, Protocol::Http1)
                                }
                            }),
                    ),
            )
            .map_err(|e| match e {
                TimeoutError::Service(e) => e,
                TimeoutError::Timeout => ConnectError::Timeout,
            });

            let tcp_service = TimeoutService::new(
                self.timeout,
                apply_fn(self.connector.clone(), |msg: Uri, srv| srv.call(msg.into()))
                    .map_err(ConnectError::from)
                    .map(|stream| (stream.into_parts().0, Protocol::Http1)),
            )
            .map_err(|e| match e {
                TimeoutError::Service(e) => e,
                TimeoutError::Timeout => ConnectError::Timeout,
            });

            connect_impl::InnerConnector {
                tcp_pool: ConnectionPool::new(
                    tcp_service,
                    self.conn_lifetime,
                    self.conn_keep_alive,
                    None,
                    self.limit,
                ),
                ssl_pool: ConnectionPool::new(
                    ssl_service,
                    self.conn_lifetime,
                    self.conn_keep_alive,
                    Some(self.disconnect_timeout),
                    self.limit,
                ),
            }
        }
    }
}

#[cfg(not(feature = "ssl"))]
mod connect_impl {
    use futures::future::{err, Either, FutureResult};
    use futures::Poll;

    use super::*;
    use crate::client::connection::IoConnection;

    pub(crate) struct InnerConnector<T, Io>
    where
        Io: AsyncRead + AsyncWrite + 'static,
        T: Service<Request = Uri, Response = (Io, Protocol), Error = ConnectError>,
    {
        pub(crate) tcp_pool: ConnectionPool<T, Io>,
    }

    impl<T, Io> Clone for InnerConnector<T, Io>
    where
        Io: AsyncRead + AsyncWrite + 'static,
        T: Service<Request = Uri, Response = (Io, Protocol), Error = ConnectError>
            + Clone,
    {
        fn clone(&self) -> Self {
            InnerConnector {
                tcp_pool: self.tcp_pool.clone(),
            }
        }
    }

    impl<T, Io> Service for InnerConnector<T, Io>
    where
        Io: AsyncRead + AsyncWrite + 'static,
        T: Service<Request = Uri, Response = (Io, Protocol), Error = ConnectError>,
    {
        type Request = Uri;
        type Response = IoConnection<Io>;
        type Error = ConnectError;
        type Future = Either<
            <ConnectionPool<T, Io> as Service>::Future,
            FutureResult<IoConnection<Io>, ConnectError>,
        >;

        fn poll_ready(&mut self) -> Poll<(), Self::Error> {
            self.tcp_pool.poll_ready()
        }

        fn call(&mut self, req: Uri) -> Self::Future {
            match req.scheme_str() {
                Some("https") | Some("wss") => {
                    Either::B(err(ConnectError::SslIsNotSupported))
                }
                _ => Either::A(self.tcp_pool.call(req)),
            }
        }
    }
}

#[cfg(feature = "ssl")]
mod connect_impl {
    use std::marker::PhantomData;

    use futures::future::{Either, FutureResult};
    use futures::{Async, Future, Poll};

    use super::*;
    use crate::client::connection::EitherConnection;

    pub(crate) struct InnerConnector<T1, T2, Io1, Io2>
    where
        Io1: AsyncRead + AsyncWrite + 'static,
        Io2: AsyncRead + AsyncWrite + 'static,
        T1: Service<Request = Uri, Response = (Io1, Protocol), Error = ConnectError>,
        T2: Service<Request = Uri, Response = (Io2, Protocol), Error = ConnectError>,
    {
        pub(crate) tcp_pool: ConnectionPool<T1, Io1>,
        pub(crate) ssl_pool: ConnectionPool<T2, Io2>,
    }

    impl<T1, T2, Io1, Io2> Clone for InnerConnector<T1, T2, Io1, Io2>
    where
        Io1: AsyncRead + AsyncWrite + 'static,
        Io2: AsyncRead + AsyncWrite + 'static,
        T1: Service<Request = Uri, Response = (Io1, Protocol), Error = ConnectError>
            + Clone,
        T2: Service<Request = Uri, Response = (Io2, Protocol), Error = ConnectError>
            + Clone,
    {
        fn clone(&self) -> Self {
            InnerConnector {
                tcp_pool: self.tcp_pool.clone(),
                ssl_pool: self.ssl_pool.clone(),
            }
        }
    }

    impl<T1, T2, Io1, Io2> Service for InnerConnector<T1, T2, Io1, Io2>
    where
        Io1: AsyncRead + AsyncWrite + 'static,
        Io2: AsyncRead + AsyncWrite + 'static,
        T1: Service<Request = Uri, Response = (Io1, Protocol), Error = ConnectError>,
        T2: Service<Request = Uri, Response = (Io2, Protocol), Error = ConnectError>,
    {
        type Request = Uri;
        type Response = EitherConnection<Io1, Io2>;
        type Error = ConnectError;
        type Future = Either<
            FutureResult<Self::Response, Self::Error>,
            Either<
                InnerConnectorResponseA<T1, Io1, Io2>,
                InnerConnectorResponseB<T2, Io1, Io2>,
            >,
        >;

        fn poll_ready(&mut self) -> Poll<(), Self::Error> {
            self.tcp_pool.poll_ready()
        }

        fn call(&mut self, req: Uri) -> Self::Future {
            match req.scheme_str() {
                Some("https") | Some("wss") => {
                    Either::B(Either::B(InnerConnectorResponseB {
                        fut: self.ssl_pool.call(req),
                        _t: PhantomData,
                    }))
                }
                _ => Either::B(Either::A(InnerConnectorResponseA {
                    fut: self.tcp_pool.call(req),
                    _t: PhantomData,
                })),
            }
        }
    }

    pub(crate) struct InnerConnectorResponseA<T, Io1, Io2>
    where
        Io1: AsyncRead + AsyncWrite + 'static,
        T: Service<Request = Uri, Response = (Io1, Protocol), Error = ConnectError>,
    {
        fut: <ConnectionPool<T, Io1> as Service>::Future,
        _t: PhantomData<Io2>,
    }

    impl<T, Io1, Io2> Future for InnerConnectorResponseA<T, Io1, Io2>
    where
        T: Service<Request = Uri, Response = (Io1, Protocol), Error = ConnectError>,
        Io1: AsyncRead + AsyncWrite + 'static,
        Io2: AsyncRead + AsyncWrite + 'static,
    {
        type Item = EitherConnection<Io1, Io2>;
        type Error = ConnectError;

        fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
            match self.fut.poll()? {
                Async::NotReady => Ok(Async::NotReady),
                Async::Ready(res) => Ok(Async::Ready(EitherConnection::A(res))),
            }
        }
    }

    pub(crate) struct InnerConnectorResponseB<T, Io1, Io2>
    where
        Io2: AsyncRead + AsyncWrite + 'static,
        T: Service<Request = Uri, Response = (Io2, Protocol), Error = ConnectError>,
    {
        fut: <ConnectionPool<T, Io2> as Service>::Future,
        _t: PhantomData<Io1>,
    }

    impl<T, Io1, Io2> Future for InnerConnectorResponseB<T, Io1, Io2>
    where
        T: Service<Request = Uri, Response = (Io2, Protocol), Error = ConnectError>,
        Io1: AsyncRead + AsyncWrite + 'static,
        Io2: AsyncRead + AsyncWrite + 'static,
    {
        type Item = EitherConnection<Io1, Io2>;
        type Error = ConnectError;

        fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
            match self.fut.poll()? {
                Async::NotReady => Ok(Async::NotReady),
                Async::Ready(res) => Ok(Async::Ready(EitherConnection::B(res))),
            }
        }
    }
}
<commit_msg>fix import<commit_after>use std::fmt;
use std::marker::PhantomData;
use std::time::Duration;

use actix_codec::{AsyncRead, AsyncWrite};
use actix_connect::{
    default_connector, Connect as TcpConnect, Connection as TcpConnection,
};
use actix_service::{apply_fn, Service, ServiceExt};
use actix_utils::timeout::{TimeoutError, TimeoutService};
use http::Uri;
use tokio_tcp::TcpStream;

use super::connection::Connection;
use super::error::ConnectError;
use super::pool::{ConnectionPool, Protocol};

#[cfg(feature = "ssl")]
use openssl::ssl::SslConnector;

#[cfg(not(feature = "ssl"))]
type SslConnector = ();

/// Http client connector builde instance.
/// `Connector` type uses builder-like pattern for connector service construction.
pub struct Connector<T, U> {
    connector: T,
    timeout: Duration,
    conn_lifetime: Duration,
    conn_keep_alive: Duration,
    disconnect_timeout: Duration,
    limit: usize,
    #[allow(dead_code)]
    ssl: SslConnector,
    _t: PhantomData<U>,
}

impl Connector<(), ()> {
    pub fn new() -> Connector<
        impl Service<
                Request = TcpConnect<Uri>,
                Response = TcpConnection<Uri, TcpStream>,
                Error = actix_connect::ConnectError,
            > + Clone,
        TcpStream,
    > {
        let ssl = {
            #[cfg(feature = "ssl")]
            {
                use log::error;
                use openssl::ssl::{SslConnector, SslMethod};

                let mut ssl = SslConnector::builder(SslMethod::tls()).unwrap();
                let _ = ssl
                    .set_alpn_protos(b"\x02h2\x08http/1.1")
                    .map_err(|e| error!("Can not set alpn protocol: {:?}", e));
                ssl.build()
            }
            #[cfg(not(feature = "ssl"))]
            {}
        };

        Connector {
            ssl,
            connector: default_connector(),
            timeout: Duration::from_secs(1),
            conn_lifetime: Duration::from_secs(75),
            conn_keep_alive: Duration::from_secs(15),
            disconnect_timeout: Duration::from_millis(3000),
            limit: 100,
            _t: PhantomData,
        }
    }
}

impl<T, U> Connector<T, U> {
    /// Use custom connector.
    pub fn connector<T1, U1>(self, connector: T1) -> Connector<T1, U1>
    where
        U1: AsyncRead + AsyncWrite + fmt::Debug,
        T1: Service<
                Request = TcpConnect<Uri>,
                Response = TcpConnection<Uri, U1>,
                Error = actix_connect::ConnectError,
            > + Clone,
    {
        Connector {
            connector,
            timeout: self.timeout,
            conn_lifetime: self.conn_lifetime,
            conn_keep_alive: self.conn_keep_alive,
            disconnect_timeout: self.disconnect_timeout,
            limit: self.limit,
            ssl: self.ssl,
            _t: PhantomData,
        }
    }
}

impl<T, U> Connector<T, U>
where
    U: AsyncRead + AsyncWrite + fmt::Debug + 'static,
    T: Service<
            Request = TcpConnect<Uri>,
            Response = TcpConnection<Uri, U>,
            Error = actix_connect::ConnectError,
        > + Clone,
{
    /// Connection timeout, i.e. max time to connect to remote host including dns name resolution.
    /// Set to 1 second by default.
    pub fn timeout(mut self, timeout: Duration) -> Self {
        self.timeout = timeout;
        self
    }

    #[cfg(feature = "ssl")]
    /// Use custom `SslConnector` instance.
    pub fn ssl(mut self, connector: SslConnector) -> Self {
        self.ssl = connector;
        self
    }

    /// Set total number of simultaneous connections per type of scheme.
    ///
    /// If limit is 0, the connector has no limit.
    /// The default limit size is 100.
    pub fn limit(mut self, limit: usize) -> Self {
        self.limit = limit;
        self
    }

    /// Set keep-alive period for opened connection.
    ///
    /// Keep-alive period is the period between connection usage. If
    /// the delay between repeated usages of the same connection
    /// exceeds this period, the connection is closed.
    /// Default keep-alive period is 15 seconds.
    pub fn conn_keep_alive(mut self, dur: Duration) -> Self {
        self.conn_keep_alive = dur;
        self
    }

    /// Set max lifetime period for connection.
    ///
    /// Connection lifetime is max lifetime of any opened connection
    /// until it is closed regardless of keep-alive period.
    /// Default lifetime period is 75 seconds.
    pub fn conn_lifetime(mut self, dur: Duration) -> Self {
        self.conn_lifetime = dur;
        self
    }

    /// Set server connection disconnect timeout in milliseconds.
    ///
    /// Defines a timeout for disconnect connection. If a disconnect procedure does not complete
    /// within this time, the socket get dropped. This timeout affects only secure connections.
    ///
    /// To disable timeout set value to 0.
    ///
    /// By default disconnect timeout is set to 3000 milliseconds.
    pub fn disconnect_timeout(mut self, dur: Duration) -> Self {
        self.disconnect_timeout = dur;
        self
    }

    /// Finish configuration process and create connector service.
    pub fn service(
        self,
    ) -> impl Service<Request = Uri, Response = impl Connection, Error = ConnectError> + Clone
    {
        #[cfg(not(feature = "ssl"))]
        {
            let connector = TimeoutService::new(
                self.timeout,
                apply_fn(self.connector, |msg: Uri, srv| srv.call(msg.into()))
                    .map_err(ConnectError::from)
                    .map(|stream| (stream.into_parts().0, Protocol::Http1)),
            )
            .map_err(|e| match e {
                TimeoutError::Service(e) => e,
                TimeoutError::Timeout => ConnectError::Timeout,
            });

            connect_impl::InnerConnector {
                tcp_pool: ConnectionPool::new(
                    connector,
                    self.conn_lifetime,
                    self.conn_keep_alive,
                    None,
                    self.limit,
                ),
            }
        }
        #[cfg(feature = "ssl")]
        {
            const H2: &[u8] = b"h2";
            use actix_connect::ssl::OpensslConnector;

            let ssl_service = TimeoutService::new(
                self.timeout,
                apply_fn(self.connector.clone(), |msg: Uri, srv| srv.call(msg.into()))
                    .map_err(ConnectError::from)
                    .and_then(
                        OpensslConnector::service(self.ssl)
                            .map_err(ConnectError::from)
                            .map(|stream| {
                                let sock = stream.into_parts().0;
                                let h2 = sock
                                    .get_ref()
                                    .ssl()
                                    .selected_alpn_protocol()
                                    .map(|protos| protos.windows(2).any(|w| w == H2))
                                    .unwrap_or(false);
                                if h2 {
                                    (sock, Protocol::Http2)
                                } else {
                                    (sock, Protocol::Http1)
                                }
                            }),
                    ),
            )
            .map_err(|e| match e {
                TimeoutError::Service(e) => e,
                TimeoutError::Timeout => ConnectError::Timeout,
            });

            let tcp_service = TimeoutService::new(
                self.timeout,
                apply_fn(self.connector.clone(), |msg: Uri, srv| srv.call(msg.into()))
                    .map_err(ConnectError::from)
                    .map(|stream| (stream.into_parts().0, Protocol::Http1)),
            )
            .map_err(|e| match e {
                TimeoutError::Service(e) => e,
                TimeoutError::Timeout => ConnectError::Timeout,
            });

            connect_impl::InnerConnector {
                tcp_pool: ConnectionPool::new(
                    tcp_service,
                    self.conn_lifetime,
                    self.conn_keep_alive,
                    None,
                    self.limit,
                ),
                ssl_pool: ConnectionPool::new(
                    ssl_service,
                    self.conn_lifetime,
                    self.conn_keep_alive,
                    Some(self.disconnect_timeout),
                    self.limit,
                ),
            }
        }
    }
}

#[cfg(not(feature = "ssl"))]
mod connect_impl {
    use futures::future::{err, Either, FutureResult};
    use futures::Poll;

    use super::*;
    use crate::client::connection::IoConnection;

    pub(crate) struct InnerConnector<T, Io>
    where
        Io: AsyncRead + AsyncWrite + 'static,
        T: Service<Request = Uri, Response = (Io, Protocol), Error = ConnectError>,
    {
        pub(crate) tcp_pool: ConnectionPool<T, Io>,
    }

    impl<T, Io> Clone for InnerConnector<T, Io>
    where
        Io: AsyncRead + AsyncWrite + 'static,
        T: Service<Request = Uri, Response = (Io, Protocol), Error = ConnectError>
            + Clone,
    {
        fn clone(&self) -> Self {
            InnerConnector {
                tcp_pool: self.tcp_pool.clone(),
            }
        }
    }

    impl<T, Io> Service for InnerConnector<T, Io>
    where
        Io: AsyncRead + AsyncWrite + 'static,
        T: Service<Request = Uri, Response = (Io, Protocol), Error = ConnectError>,
    {
        type Request = Uri;
        type Response = IoConnection<Io>;
        type Error = ConnectError;
        type Future = Either<
            <ConnectionPool<T, Io> as Service>::Future,
            FutureResult<IoConnection<Io>, ConnectError>,
        >;

        fn poll_ready(&mut self) -> Poll<(), Self::Error> {
            self.tcp_pool.poll_ready()
        }

        fn call(&mut self, req: Uri) -> Self::Future {
            match req.scheme_str() {
                Some("https") | Some("wss") => {
                    Either::B(err(ConnectError::SslIsNotSupported))
                }
                _ => Either::A(self.tcp_pool.call(req)),
            }
        }
    }
}

#[cfg(feature = "ssl")]
mod connect_impl {
    use std::marker::PhantomData;

    use futures::future::{Either, FutureResult};
    use futures::{Async, Future, Poll};

    use super::*;
    use crate::client::connection::EitherConnection;

    pub(crate) struct InnerConnector<T1, T2, Io1, Io2>
    where
        Io1: AsyncRead + AsyncWrite + 'static,
        Io2: AsyncRead + AsyncWrite + 'static,
        T1: Service<Request = Uri, Response = (Io1, Protocol), Error = ConnectError>,
        T2: Service<Request = Uri, Response = (Io2, Protocol), Error = ConnectError>,
    {
        pub(crate) tcp_pool: ConnectionPool<T1, Io1>,
        pub(crate) ssl_pool: ConnectionPool<T2, Io2>,
    }

    impl<T1, T2, Io1, Io2> Clone for InnerConnector<T1, T2, Io1, Io2>
    where
        Io1: AsyncRead + AsyncWrite + 'static,
        Io2: AsyncRead + AsyncWrite + 'static,
        T1: Service<Request = Uri, Response = (Io1, Protocol), Error = ConnectError>
            + Clone,
        T2: Service<Request = Uri, Response = (Io2, Protocol), Error = ConnectError>
            + Clone,
    {
        fn clone(&self) -> Self {
            InnerConnector {
                tcp_pool: self.tcp_pool.clone(),
                ssl_pool: self.ssl_pool.clone(),
            }
        }
    }

    impl<T1, T2, Io1, Io2> Service for InnerConnector<T1, T2, Io1, Io2>
    where
        Io1: AsyncRead + AsyncWrite + 'static,
        Io2: AsyncRead + AsyncWrite + 'static,
        T1: Service<Request = Uri, Response = (Io1, Protocol), Error = ConnectError>,
        T2: Service<Request = Uri, Response = (Io2, Protocol), Error = ConnectError>,
    {
        type Request = Uri;
        type Response = EitherConnection<Io1, Io2>;
        type Error = ConnectError;
        type Future = Either<
            FutureResult<Self::Response, Self::Error>,
            Either<
                InnerConnectorResponseA<T1, Io1, Io2>,
                InnerConnectorResponseB<T2, Io1, Io2>,
            >,
        >;

        fn poll_ready(&mut self) -> Poll<(), Self::Error> {
            self.tcp_pool.poll_ready()
        }

        fn call(&mut self, req: Uri) -> Self::Future {
            match req.scheme_str() {
                Some("https") | Some("wss") => {
                    Either::B(Either::B(InnerConnectorResponseB {
                        fut: self.ssl_pool.call(req),
                        _t: PhantomData,
                    }))
                }
                _ => Either::B(Either::A(InnerConnectorResponseA {
                    fut: self.tcp_pool.call(req),
                    _t: PhantomData,
                })),
            }
        }
    }

    pub(crate) struct InnerConnectorResponseA<T, Io1, Io2>
    where
        Io1: AsyncRead + AsyncWrite + 'static,
        T: Service<Request = Uri, Response = (Io1, Protocol), Error = ConnectError>,
    {
        fut: <ConnectionPool<T, Io1> as Service>::Future,
        _t: PhantomData<Io2>,
    }

    impl<T, Io1, Io2> Future for InnerConnectorResponseA<T, Io1, Io2>
    where
        T: Service<Request = Uri, Response = (Io1, Protocol), Error = ConnectError>,
        Io1: AsyncRead + AsyncWrite + 'static,
        Io2: AsyncRead + AsyncWrite + 'static,
    {
        type Item = EitherConnection<Io1, Io2>;
        type Error = ConnectError;

        fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
            match self.fut.poll()? {
                Async::NotReady => Ok(Async::NotReady),
                Async::Ready(res) => Ok(Async::Ready(EitherConnection::A(res))),
            }
        }
    }

    pub(crate) struct InnerConnectorResponseB<T, Io1, Io2>
    where
        Io2: AsyncRead + AsyncWrite + 'static,
        T: Service<Request = Uri, Response = (Io2, Protocol), Error = ConnectError>,
    {
        fut: <ConnectionPool<T, Io2> as Service>::Future,
        _t: PhantomData<Io1>,
    }

    impl<T, Io1, Io2> Future for InnerConnectorResponseB<T, Io1, Io2>
    where
        T: Service<Request = Uri, Response = (Io2, Protocol), Error = ConnectError>,
        Io1: AsyncRead + AsyncWrite + 'static,
        Io2: AsyncRead + AsyncWrite + 'static,
    {
        type Item = EitherConnection<Io1, Io2>;
        type Error = ConnectError;

        fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
            match self.fut.poll()? {
                Async::NotReady => Ok(Async::NotReady),
                Async::Ready(res) => Ok(Async::Ready(EitherConnection::B(res))),
            }
        }
    }
}
<|endoftext|>