<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

//! Types and methods for looking up locations.
//!
//! Location search for Twitter works in one of two ways. The most direct method is to take a
//! latitude/longitude coordinate (say, from a devide's GPS system or by geolocating from wi-fi
//! networks, or simply from a known coordinate) and call `reverse_geocode`. Twitter says
//! `reverse_geocode` provides more of a "raw data access", and it can be considered to merely show
//! what locations are in that point or area.
//!
//! On the other hand, if you're intending to let a user select from a list of locations, you can
//! use the `search_*` methods instead. These have much of the same available parameters, but will
//! "potentially re-order [results] with regards to the user who is authenticated." In addition,
//! the results may potentially pull in "nearby" results to allow for a more broad selection or to
//! account for inaccurate location reporting.
//!
//! Since there are several optional parameters to both query methods, each one is assembled as a
//! builder. You can create the builder with the `reverse_geocode`, `search_point`, `search_query`,
//! or `search_ip` functions. From there, add any additional parameters by chaining method calls
//! onto the builder. When you're ready to peform the search call, hand your tokens to `call`, and
//! the list of results will be returned.
//!
//! Along with the list of place results, Twitter also returns the full search URL. egg-mode
//! returns this URL as part of the result struct, allowing you to perform the same search using
//! the `reverse_geocode_url` or `search_url` functions.

use std::collections::HashMap;
use std::fmt;

use serde::{Deserialize, Deserializer};
use serde::de::Error;
use serde_json;

use auth;
use common::*;
use links;

mod fun;

pub use self::fun::*;

// https://developer.twitter.com/en/docs/tweets/data-dictionary/overview/geo-objects#place
///Represents a named location.
#[derive(Debug, Clone, Deserialize)]
pub struct Place {
    ///Alphanumeric ID of the location.
    pub id: String,
    ///Map of miscellaneous information about this place. See [Twitter's documentation][attrib] for
    ///details and common attribute keys.
    ///
    ///[attrib]: https://dev.twitter.com/overview/api/places#attributes
    pub attributes: HashMap<String, String>,
    ///A bounding box of latitude/longitude coordinates that encloses this place.
    #[serde(deserialize_with = "deserialize_bounding_box")]
    pub bounding_box: Vec<(f64, f64)>,
    ///Name of the country containing this place.
    pub country: String,
    ///Shortened country code representing the country containing this place.
    pub country_code: String,
    ///Full human-readable name of this place.
    pub full_name: String,
    ///Short human-readable name of this place.
    pub name: String,
    ///The type of location represented by this place.
    pub place_type: PlaceType,
    ///If present, the country or administrative region that contains this place.
    pub contained_within: Option<Vec<Place>>,
}

///Represents the type of region represented by a given place.
#[derive(Debug, Copy, Clone, Serialize, Deserialize)]
pub enum PlaceType {
    ///A coordinate with no area.
    #[serde(rename = "poi")]
    PointOfInterest,
    ///A region within a city.
    #[serde(rename = "neighborhood")]
    Neighborhood,
    ///An entire city.
    #[serde(rename = "city")]
    City,
    ///An administrative area, e.g. state or province.
    #[serde(rename = "admin")]
    Admin,
    ///An entire country.
    #[serde(rename = "country")]
    Country,
}

///Represents the accuracy of a GPS measurement, when being given to a location search.
#[derive(Debug, Copy, Clone)]
pub enum Accuracy {
    ///Location accurate to the given number of meters.
    Meters(f64),
    ///Location accurate to the given number of feet.
    Feet(f64),
}

///Represents the result of a location search, either via `reverse_geocode` or `search`.
pub struct SearchResult {
    ///The full URL used to pull the result list. This can be fed to the `_url` version of your
    ///original call to avoid having to fill out the argument list again.
    pub url: String,
    ///The list of results from the search.
    pub results: Vec<Place>,
}

impl<'de> Deserialize<'de> for SearchResult {
    fn deserialize<D>(deser: D) -> Result<SearchResult, D::Error> where D: Deserializer<'de> {
        let raw: serde_json::Value = serde_json::Value::deserialize(deser)?;
        let url = raw.get("query")
            .and_then(|obj| obj.get("url"))
            .ok_or_else(|| D::Error::custom("Malformed search result"))?
            .to_string();
        let results = raw.get("result")
            .and_then(|obj| obj.get("places"))
            .and_then(|arr| <Vec<Place>>::deserialize(arr).ok())
            .ok_or_else(|| D::Error::custom("Malformed search result"))?;
        Ok(SearchResult {
            url, results
        })
    }
}

///Represents a `reverse_geocode` query before it is sent.
///
///The available methods on this builder struct allow you to specify optional parameters to the
///search operation. Where applicable, each method lists its default value and acceptable ranges.
///
///To complete your search setup and send the query to Twitter, hand your tokens to `call`. The
///list of results from Twitter will be returned, as well as a URL to perform the same search via
///`reverse_geocode_url`.
pub struct GeocodeBuilder {
    coordinate: (f64, f64),
    accuracy: Option<Accuracy>,
    granularity: Option<PlaceType>,
    max_results: Option<u32>,
}

impl GeocodeBuilder {
    ///Begins building a reverse-geocode query with the given coordinate.
    fn new(latitude: f64, longitude: f64) -> Self {
        GeocodeBuilder {
            coordinate: (latitude, longitude),
            accuracy: None,
            granularity: None,
            max_results: None,
        }
    }

    ///Expands the area to search to the given radius. By default, this is zero.
    ///
    ///From Twitter: "If coming from a device, in practice, this value is whatever accuracy the
    ///device has measuring its location (whether it be coming from a GPS, WiFi triangulation,
    ///etc.)."
    pub fn accuracy(self, accuracy: Accuracy) -> Self {
        GeocodeBuilder {
            accuracy: Some(accuracy),
            ..self
        }
    }

    ///Sets the minimal specificity of what kind of results to return. For example, passing `City`
    ///to this will make the eventual result exclude neighborhoods and points.
    pub fn granularity(self, granularity: PlaceType) -> Self {
        GeocodeBuilder {
            granularity: Some(granularity),
            ..self
        }
    }

    ///Restricts the maximum number of results returned in this search. This is not a guarantee
    ///that the search will return this many results, but instead provides a hint as to how many
    ///"nearby" results to return.
    ///
    ///This value has a default value of 20, which is also its maximum. If zero or a number greater
    ///than 20 is passed here, it will be defaulted to 20 before sending to Twitter.
    ///
    ///From Twitter: "Ideally, only pass in the number of places you intend to display to the user
    ///here."
    pub fn max_results(self, max_results: u32) -> Self {
        GeocodeBuilder {
            max_results: Some(max_results),
            ..self
        }
    }

    ///Finalize the search parameters and return the results collection.
    pub fn call(&self, token: &auth::Token)
        -> FutureResponse<SearchResult>
    {
        let mut params = HashMap::new();

        add_param(&mut params, "lat", self.coordinate.0.to_string());
        add_param(&mut params, "long", self.coordinate.1.to_string());

        if let Some(ref accuracy) = self.accuracy {
            add_param(&mut params, "accuracy", accuracy.to_string());
        }

        if let Some(ref param) = self.granularity {
            add_param(&mut params, "granularity", param.to_string());
        }

        if let Some(count) = self.max_results {
            let count = if count == 0 || count > 20 { 20 } else { count };
            add_param(&mut params, "max_results", count.to_string());
        }

        let req = auth::get(links::place::REVERSE_GEOCODE, token, Some(&params));

        make_parsed_future(req)
    }
}

enum PlaceQuery<'a> {
    LatLon(f64, f64),
    Query(&'a str),
    IPAddress(&'a str),
}

///Represents a location search query before it is sent.
///
///The available methods on this builder struct allow you to specify optional parameters to the
///search operation. Where applicable, each method lists its default value and acceptable ranges.
///
///To complete your search setup and send the query to Twitter, hand your tokens to `call`. The
///list of results from Twitter will be returned, as well as a URL to perform the same search via
///`search_url`.
pub struct SearchBuilder<'a> {
    query: PlaceQuery<'a>,
    accuracy: Option<Accuracy>,
    granularity: Option<PlaceType>,
    max_results: Option<u32>,
    contained_within: Option<&'a str>,
    attributes: Option<HashMap<&'a str, &'a str>>,
}

impl<'a> SearchBuilder<'a> {
    ///Begins building a location search with the given query.
    fn new(query: PlaceQuery<'a>) -> Self {
        SearchBuilder {
            query: query,
            accuracy: None,
            granularity: None,
            max_results: None,
            contained_within: None,
            attributes: None,
        }
    }

    ///Expands the area to search to the given radius. By default, this is zero.
    ///
    ///From Twitter: "If coming from a device, in practice, this value is whatever accuracy the
    ///device has measuring its location (whether it be coming from a GPS, WiFi triangulation,
    ///etc.)."
    pub fn accuracy(self, accuracy: Accuracy) -> Self {
        SearchBuilder {
            accuracy: Some(accuracy),
            ..self
        }
    }

    ///Sets the minimal specificity of what kind of results to return. For example, passing `City`
    ///to this will make the eventual result exclude neighborhoods and points.
    pub fn granularity(self, granularity: PlaceType) -> Self {
        SearchBuilder {
            granularity: Some(granularity),
            ..self
        }
    }

    ///Restricts the maximum number of results returned in this search. This is not a guarantee
    ///that the search will return this many results, but instead provides a hint as to how many
    ///"nearby" results to return.
    ///
    ///From experimentation, this value has a default of 20 and a maximum of 100. If fewer
    ///locations match the search parameters, fewer places will be returned.
    ///
    ///From Twitter: "Ideally, only pass in the number of places you intend to display to the user
    ///here."
    pub fn max_results(self, max_results: u32) -> Self {
        SearchBuilder {
            max_results: Some(max_results),
            ..self
        }
    }

    ///Restricts results to those contained within the given Place ID.
    pub fn contained_within(self, contained_id: &'a str) -> Self {
        SearchBuilder {
            contained_within: Some(contained_id),
            ..self
        }
    }

    ///Restricts results to those with the given attribute. A list of common attributes are
    ///available in [Twitter's documentation for Places][attrs]. Custom attributes are supported in
    ///this search, if you know them. This function may be called multiple times with different
    ///`attribute_key` values to combine attribute search parameters.
    ///
    ///[attrs]: https://dev.twitter.com/overview/api/places#attributes
    ///
    ///For example, `.attribute("street_address", "123 Main St")` searches for places with the
    ///given street address.
    pub fn attribute(self, attribute_key: &'a str, attribute_value: &'a str) -> Self {
        let mut attrs = self.attributes.unwrap_or_default();
        attrs.insert(attribute_key, attribute_value);

        SearchBuilder {
            attributes: Some(attrs),
            ..self
        }
    }

    ///Finalize the search parameters and return the results collection.
    pub fn call(&self, token: &auth::Token)
        -> FutureResponse<SearchResult>
    {
        let mut params = HashMap::new();

        match self.query {
            PlaceQuery::LatLon(lat, long) => {
                add_param(&mut params, "lat", lat.to_string());
                add_param(&mut params, "long", long.to_string());
            },
            PlaceQuery::Query(text) => {
                add_param(&mut params, "query", text);
            },
            PlaceQuery::IPAddress(text) => {
                add_param(&mut params, "ip", text);
            },
        }

        if let Some(ref acc) = self.accuracy {
            add_param(&mut params, "accuracy", acc.to_string());
        }

        if let Some(ref gran) = self.granularity {
            add_param(&mut params, "granularity", gran.to_string());
        }

        if let Some(max) = self.max_results {
            add_param(&mut params, "max_results", max.to_string());
        }

        if let Some(id) = self.contained_within {
            add_param(&mut params, "contained_within", id);
        }

        if let Some(ref attrs) = self.attributes {
            for (k, v) in attrs {
                add_param(&mut params, format!("attribute:{}", k), *v);
            }
        }

        let req = auth::get(links::place::SEARCH, token, Some(&params));

        make_parsed_future(req)
    }
}

///Display impl to make `to_string()` format the enum for sending to Twitter. This is *mostly* just
///a lowercase version of the variants, but `Point` is rendered as `"poi"` instead.
impl fmt::Display for PlaceType {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let quoted = serde_json::to_string(self).unwrap();
        let inner = &quoted[1..quoted.len() - 1]; // ignore the quote marks
        write!(f, "{}", inner)
    }
}

///Display impl to make `to_string()` format the enum for sending to Twitter. This turns `Meters`
///into the contained number by itself, and `Feet` into the number suffixed by `"ft"`.
impl fmt::Display for Accuracy {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Accuracy::Meters(dist) => write!(f, "{}", dist),
            Accuracy::Feet(dist) => write!(f, "{}ft", dist),
        }
    }
}

fn deserialize_bounding_box<'de, D>(ser: D) -> Result<Vec<(f64, f64)>, D::Error> where D: Deserializer<'de> {
    let s = serde_json::Value::deserialize(ser)?;
    s.get("coordinates")
        .and_then(|arr| arr.get(0).cloned())
        .ok_or_else(|| D::Error::custom("Malformed 'bounding_box' attribute"))
        .and_then(|inner_arr| serde_json::from_value::<Vec<(f64, f64)>>(inner_arr)
                .map_err(|e| D::Error::custom(e))
        )
}
<commit_msg>escape square brackets to silence intra-doc link warning<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

//! Types and methods for looking up locations.
//!
//! Location search for Twitter works in one of two ways. The most direct method is to take a
//! latitude/longitude coordinate (say, from a devide's GPS system or by geolocating from wi-fi
//! networks, or simply from a known coordinate) and call `reverse_geocode`. Twitter says
//! `reverse_geocode` provides more of a "raw data access", and it can be considered to merely show
//! what locations are in that point or area.
//!
//! On the other hand, if you're intending to let a user select from a list of locations, you can
//! use the `search_*` methods instead. These have much of the same available parameters, but will
//! "potentially re-order \[results\] with regards to the user who is authenticated." In addition,
//! the results may potentially pull in "nearby" results to allow for a more broad selection or to
//! account for inaccurate location reporting.
//!
//! Since there are several optional parameters to both query methods, each one is assembled as a
//! builder. You can create the builder with the `reverse_geocode`, `search_point`, `search_query`,
//! or `search_ip` functions. From there, add any additional parameters by chaining method calls
//! onto the builder. When you're ready to peform the search call, hand your tokens to `call`, and
//! the list of results will be returned.
//!
//! Along with the list of place results, Twitter also returns the full search URL. egg-mode
//! returns this URL as part of the result struct, allowing you to perform the same search using
//! the `reverse_geocode_url` or `search_url` functions.

use std::collections::HashMap;
use std::fmt;

use serde::{Deserialize, Deserializer};
use serde::de::Error;
use serde_json;

use auth;
use common::*;
use links;

mod fun;

pub use self::fun::*;

// https://developer.twitter.com/en/docs/tweets/data-dictionary/overview/geo-objects#place
///Represents a named location.
#[derive(Debug, Clone, Deserialize)]
pub struct Place {
    ///Alphanumeric ID of the location.
    pub id: String,
    ///Map of miscellaneous information about this place. See [Twitter's documentation][attrib] for
    ///details and common attribute keys.
    ///
    ///[attrib]: https://dev.twitter.com/overview/api/places#attributes
    pub attributes: HashMap<String, String>,
    ///A bounding box of latitude/longitude coordinates that encloses this place.
    #[serde(deserialize_with = "deserialize_bounding_box")]
    pub bounding_box: Vec<(f64, f64)>,
    ///Name of the country containing this place.
    pub country: String,
    ///Shortened country code representing the country containing this place.
    pub country_code: String,
    ///Full human-readable name of this place.
    pub full_name: String,
    ///Short human-readable name of this place.
    pub name: String,
    ///The type of location represented by this place.
    pub place_type: PlaceType,
    ///If present, the country or administrative region that contains this place.
    pub contained_within: Option<Vec<Place>>,
}

///Represents the type of region represented by a given place.
#[derive(Debug, Copy, Clone, Serialize, Deserialize)]
pub enum PlaceType {
    ///A coordinate with no area.
    #[serde(rename = "poi")]
    PointOfInterest,
    ///A region within a city.
    #[serde(rename = "neighborhood")]
    Neighborhood,
    ///An entire city.
    #[serde(rename = "city")]
    City,
    ///An administrative area, e.g. state or province.
    #[serde(rename = "admin")]
    Admin,
    ///An entire country.
    #[serde(rename = "country")]
    Country,
}

///Represents the accuracy of a GPS measurement, when being given to a location search.
#[derive(Debug, Copy, Clone)]
pub enum Accuracy {
    ///Location accurate to the given number of meters.
    Meters(f64),
    ///Location accurate to the given number of feet.
    Feet(f64),
}

///Represents the result of a location search, either via `reverse_geocode` or `search`.
pub struct SearchResult {
    ///The full URL used to pull the result list. This can be fed to the `_url` version of your
    ///original call to avoid having to fill out the argument list again.
    pub url: String,
    ///The list of results from the search.
    pub results: Vec<Place>,
}

impl<'de> Deserialize<'de> for SearchResult {
    fn deserialize<D>(deser: D) -> Result<SearchResult, D::Error> where D: Deserializer<'de> {
        let raw: serde_json::Value = serde_json::Value::deserialize(deser)?;
        let url = raw.get("query")
            .and_then(|obj| obj.get("url"))
            .ok_or_else(|| D::Error::custom("Malformed search result"))?
            .to_string();
        let results = raw.get("result")
            .and_then(|obj| obj.get("places"))
            .and_then(|arr| <Vec<Place>>::deserialize(arr).ok())
            .ok_or_else(|| D::Error::custom("Malformed search result"))?;
        Ok(SearchResult {
            url, results
        })
    }
}

///Represents a `reverse_geocode` query before it is sent.
///
///The available methods on this builder struct allow you to specify optional parameters to the
///search operation. Where applicable, each method lists its default value and acceptable ranges.
///
///To complete your search setup and send the query to Twitter, hand your tokens to `call`. The
///list of results from Twitter will be returned, as well as a URL to perform the same search via
///`reverse_geocode_url`.
pub struct GeocodeBuilder {
    coordinate: (f64, f64),
    accuracy: Option<Accuracy>,
    granularity: Option<PlaceType>,
    max_results: Option<u32>,
}

impl GeocodeBuilder {
    ///Begins building a reverse-geocode query with the given coordinate.
    fn new(latitude: f64, longitude: f64) -> Self {
        GeocodeBuilder {
            coordinate: (latitude, longitude),
            accuracy: None,
            granularity: None,
            max_results: None,
        }
    }

    ///Expands the area to search to the given radius. By default, this is zero.
    ///
    ///From Twitter: "If coming from a device, in practice, this value is whatever accuracy the
    ///device has measuring its location (whether it be coming from a GPS, WiFi triangulation,
    ///etc.)."
    pub fn accuracy(self, accuracy: Accuracy) -> Self {
        GeocodeBuilder {
            accuracy: Some(accuracy),
            ..self
        }
    }

    ///Sets the minimal specificity of what kind of results to return. For example, passing `City`
    ///to this will make the eventual result exclude neighborhoods and points.
    pub fn granularity(self, granularity: PlaceType) -> Self {
        GeocodeBuilder {
            granularity: Some(granularity),
            ..self
        }
    }

    ///Restricts the maximum number of results returned in this search. This is not a guarantee
    ///that the search will return this many results, but instead provides a hint as to how many
    ///"nearby" results to return.
    ///
    ///This value has a default value of 20, which is also its maximum. If zero or a number greater
    ///than 20 is passed here, it will be defaulted to 20 before sending to Twitter.
    ///
    ///From Twitter: "Ideally, only pass in the number of places you intend to display to the user
    ///here."
    pub fn max_results(self, max_results: u32) -> Self {
        GeocodeBuilder {
            max_results: Some(max_results),
            ..self
        }
    }

    ///Finalize the search parameters and return the results collection.
    pub fn call(&self, token: &auth::Token)
        -> FutureResponse<SearchResult>
    {
        let mut params = HashMap::new();

        add_param(&mut params, "lat", self.coordinate.0.to_string());
        add_param(&mut params, "long", self.coordinate.1.to_string());

        if let Some(ref accuracy) = self.accuracy {
            add_param(&mut params, "accuracy", accuracy.to_string());
        }

        if let Some(ref param) = self.granularity {
            add_param(&mut params, "granularity", param.to_string());
        }

        if let Some(count) = self.max_results {
            let count = if count == 0 || count > 20 { 20 } else { count };
            add_param(&mut params, "max_results", count.to_string());
        }

        let req = auth::get(links::place::REVERSE_GEOCODE, token, Some(&params));

        make_parsed_future(req)
    }
}

enum PlaceQuery<'a> {
    LatLon(f64, f64),
    Query(&'a str),
    IPAddress(&'a str),
}

///Represents a location search query before it is sent.
///
///The available methods on this builder struct allow you to specify optional parameters to the
///search operation. Where applicable, each method lists its default value and acceptable ranges.
///
///To complete your search setup and send the query to Twitter, hand your tokens to `call`. The
///list of results from Twitter will be returned, as well as a URL to perform the same search via
///`search_url`.
pub struct SearchBuilder<'a> {
    query: PlaceQuery<'a>,
    accuracy: Option<Accuracy>,
    granularity: Option<PlaceType>,
    max_results: Option<u32>,
    contained_within: Option<&'a str>,
    attributes: Option<HashMap<&'a str, &'a str>>,
}

impl<'a> SearchBuilder<'a> {
    ///Begins building a location search with the given query.
    fn new(query: PlaceQuery<'a>) -> Self {
        SearchBuilder {
            query: query,
            accuracy: None,
            granularity: None,
            max_results: None,
            contained_within: None,
            attributes: None,
        }
    }

    ///Expands the area to search to the given radius. By default, this is zero.
    ///
    ///From Twitter: "If coming from a device, in practice, this value is whatever accuracy the
    ///device has measuring its location (whether it be coming from a GPS, WiFi triangulation,
    ///etc.)."
    pub fn accuracy(self, accuracy: Accuracy) -> Self {
        SearchBuilder {
            accuracy: Some(accuracy),
            ..self
        }
    }

    ///Sets the minimal specificity of what kind of results to return. For example, passing `City`
    ///to this will make the eventual result exclude neighborhoods and points.
    pub fn granularity(self, granularity: PlaceType) -> Self {
        SearchBuilder {
            granularity: Some(granularity),
            ..self
        }
    }

    ///Restricts the maximum number of results returned in this search. This is not a guarantee
    ///that the search will return this many results, but instead provides a hint as to how many
    ///"nearby" results to return.
    ///
    ///From experimentation, this value has a default of 20 and a maximum of 100. If fewer
    ///locations match the search parameters, fewer places will be returned.
    ///
    ///From Twitter: "Ideally, only pass in the number of places you intend to display to the user
    ///here."
    pub fn max_results(self, max_results: u32) -> Self {
        SearchBuilder {
            max_results: Some(max_results),
            ..self
        }
    }

    ///Restricts results to those contained within the given Place ID.
    pub fn contained_within(self, contained_id: &'a str) -> Self {
        SearchBuilder {
            contained_within: Some(contained_id),
            ..self
        }
    }

    ///Restricts results to those with the given attribute. A list of common attributes are
    ///available in [Twitter's documentation for Places][attrs]. Custom attributes are supported in
    ///this search, if you know them. This function may be called multiple times with different
    ///`attribute_key` values to combine attribute search parameters.
    ///
    ///[attrs]: https://dev.twitter.com/overview/api/places#attributes
    ///
    ///For example, `.attribute("street_address", "123 Main St")` searches for places with the
    ///given street address.
    pub fn attribute(self, attribute_key: &'a str, attribute_value: &'a str) -> Self {
        let mut attrs = self.attributes.unwrap_or_default();
        attrs.insert(attribute_key, attribute_value);

        SearchBuilder {
            attributes: Some(attrs),
            ..self
        }
    }

    ///Finalize the search parameters and return the results collection.
    pub fn call(&self, token: &auth::Token)
        -> FutureResponse<SearchResult>
    {
        let mut params = HashMap::new();

        match self.query {
            PlaceQuery::LatLon(lat, long) => {
                add_param(&mut params, "lat", lat.to_string());
                add_param(&mut params, "long", long.to_string());
            },
            PlaceQuery::Query(text) => {
                add_param(&mut params, "query", text);
            },
            PlaceQuery::IPAddress(text) => {
                add_param(&mut params, "ip", text);
            },
        }

        if let Some(ref acc) = self.accuracy {
            add_param(&mut params, "accuracy", acc.to_string());
        }

        if let Some(ref gran) = self.granularity {
            add_param(&mut params, "granularity", gran.to_string());
        }

        if let Some(max) = self.max_results {
            add_param(&mut params, "max_results", max.to_string());
        }

        if let Some(id) = self.contained_within {
            add_param(&mut params, "contained_within", id);
        }

        if let Some(ref attrs) = self.attributes {
            for (k, v) in attrs {
                add_param(&mut params, format!("attribute:{}", k), *v);
            }
        }

        let req = auth::get(links::place::SEARCH, token, Some(&params));

        make_parsed_future(req)
    }
}

///Display impl to make `to_string()` format the enum for sending to Twitter. This is *mostly* just
///a lowercase version of the variants, but `Point` is rendered as `"poi"` instead.
impl fmt::Display for PlaceType {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let quoted = serde_json::to_string(self).unwrap();
        let inner = &quoted[1..quoted.len() - 1]; // ignore the quote marks
        write!(f, "{}", inner)
    }
}

///Display impl to make `to_string()` format the enum for sending to Twitter. This turns `Meters`
///into the contained number by itself, and `Feet` into the number suffixed by `"ft"`.
impl fmt::Display for Accuracy {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Accuracy::Meters(dist) => write!(f, "{}", dist),
            Accuracy::Feet(dist) => write!(f, "{}ft", dist),
        }
    }
}

fn deserialize_bounding_box<'de, D>(ser: D) -> Result<Vec<(f64, f64)>, D::Error> where D: Deserializer<'de> {
    let s = serde_json::Value::deserialize(ser)?;
    s.get("coordinates")
        .and_then(|arr| arr.get(0).cloned())
        .ok_or_else(|| D::Error::custom("Malformed 'bounding_box' attribute"))
        .and_then(|inner_arr| serde_json::from_value::<Vec<(f64, f64)>>(inner_arr)
                .map_err(|e| D::Error::custom(e))
        )
}
<|endoftext|>


<commit_before><commit_msg>Crop uploads to texture size.<commit_after><|endoftext|>


<commit_before>use ::cairo;
use ::glib::translate::*;
use ::glib_sys;
use ::libc;

#[cfg(test)]
use std::f64::consts::*;

use cairo::MatrixTrait;
use cssparser::{self, Parser, ParserInput};

use error::*;
use parsers::{ParseError, Parse, optional_comma};

impl Parse for cairo::Matrix {
    type Data = ();
    type Err = AttributeError;

    fn parse (s: &str, _: ()) -> Result<cairo::Matrix, AttributeError> {
        parse_transform (s)
    }
}

pub fn parse_transform(s: &str) -> Result<cairo::Matrix, AttributeError> {
    let mut input = ParserInput::new(s);
    let mut parser = Parser::new(&mut input);

    let xform = parser.expect_ident_cloned()?;
    let _ = parser.expect_parenthesis_block()?;

    match xform.as_ref() {
        "matrix"    => parse_matrix_args(&mut parser),
        "translate" => parse_translate_args(&mut parser),
        "scale"     => parse_scale_args(&mut parser),
        "rotate"    => parse_rotate_args(&mut parser),
        "skewX"     => parse_skewX_args(&mut parser),
        "skewY"     => parse_skewY_args(&mut parser),
        _           => Err(AttributeError::from(ParseError::new("expected matrix|translate|scale|rotate|skewX|skewY"))),
    }
}

fn parse_matrix_args(parser: &mut Parser) -> Result<cairo::Matrix, AttributeError> {
    parser.parse_nested_block(|p| {
        let xx = p.expect_number()? as f64;
        optional_comma(p);

        let yx = p.expect_number()? as f64;
        optional_comma(p);

        let xy = p.expect_number()? as f64;
        optional_comma(p);

        let yy = p.expect_number()? as f64;
        optional_comma(p);

        let x0 = p.expect_number()? as f64;
        optional_comma(p);

        let y0 = p.expect_number()? as f64;

        Ok(cairo::Matrix::new(xx, yx, xy, yy, x0, y0))
    }).map_err(cssparser::ParseError::<()>::basic)
        .map_err(|e| AttributeError::from(e))
}

fn parse_translate_args(parser: &mut Parser) -> Result<cairo::Matrix, AttributeError> {
    unimplemented!();
}

fn parse_scale_args(parser: &mut Parser) -> Result<cairo::Matrix, AttributeError> {
    unimplemented!();
}

fn parse_rotate_args(parser: &mut Parser) -> Result<cairo::Matrix, AttributeError> {
    unimplemented!();
}

fn parse_skewX_args(parser: &mut Parser) -> Result<cairo::Matrix, AttributeError> {
    unimplemented!();
}

fn parse_skewY_args(parser: &mut Parser) -> Result<cairo::Matrix, AttributeError> {
    unimplemented!();
}

/*
pub fn parse_transform (s: &str) -> Result <cairo::Matrix, AttributeError> {
    let r = parse_TransformList (s);

    match r {
        Ok (m) => {
            m.try_invert ().map (|_| m)
                .map_err (|_| AttributeError::Value ("invalid transformation matrix".to_string ()))
        },

        Err (e) => {
            Err (AttributeError::Parse (ParseError::new (format! ("{:?}", e))))
        }
    }
}
*/

#[cfg(test)]
fn make_rotation_matrix (angle_degrees: f64, tx: f64, ty: f64) -> cairo::Matrix {
    let angle = angle_degrees * PI / 180.0;

    let mut m = cairo::Matrix::new (1.0, 0.0, 0.0, 1.0, tx, ty);

    let mut r = cairo::Matrix::identity ();
    r.rotate (angle);
    m = cairo::Matrix::multiply (&r, &m);

    m = cairo::Matrix::multiply (&cairo::Matrix::new (1.0, 0.0, 0.0, 1.0, -tx, -ty), &m);
    m
}

#[no_mangle]
pub extern fn rsvg_parse_transform (out_matrix: *mut cairo::Matrix, s: *const libc::c_char) -> glib_sys::gboolean {
    assert! (!out_matrix.is_null ());
    assert! (!s.is_null ());

    let string = unsafe { String::from_glib_none (s) };
    let matrix: &mut cairo::Matrix = unsafe { &mut *out_matrix };

    match parse_transform (&string) {
        Ok (m) => {
            *matrix = m;
            true.to_glib ()
        },

        Err (_) => {
            *matrix = cairo::Matrix::identity ();
            false.to_glib ()
        }
    }
}

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

    #[test]
    #[ignore]
    fn parses_valid_transform () {
        let t = cairo::Matrix::new (1.0, 0.0, 0.0, 1.0, 20.0, 30.0);
        let s = cairo::Matrix::new (10.0, 0.0, 0.0, 10.0, 0.0, 0.0);
        let r = make_rotation_matrix (30.0, 10.0, 10.0);

        let a = cairo::Matrix::multiply (&s, &t);
        assert_eq! (parse_transform ("translate(20, 30), scale (10) rotate (30 10 10)").unwrap (),
                    cairo::Matrix::multiply (&r, &a));
    }

    #[test]
    #[ignore]
    fn syntax_error_yields_parse_error () {
        match parse_transform ("foo") {
            Err (AttributeError::Parse (_)) => {},
            _ => { panic! (); }
        }

        match parse_transform ("matrix (1 2 3 4 5)") {
            Err (AttributeError::Parse (_)) => {},
            _ => { panic! (); }
        }
    }

    #[test]
    #[ignore]
    fn invalid_transform_yields_value_error () {
        match parse_transform ("matrix (0 0 0 0 0 0)") {
            Err (AttributeError::Value (_)) => {},
            _ => { panic! (); }
        }

        match parse_transform ("scale (0), translate (10, 10)") {
            Err (AttributeError::Value (_)) => {},
            _ => { panic! (); }
        }

        match parse_transform ("scale (0), skewX (90)") {
            Err (AttributeError::Value (_)) => {},
            _ => { panic! (); }
        }
    }
}

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

    #[test]
    fn parses_matrix () {
        assert_eq! (parse_transform ("matrix (1 2 3 4 5 6)").unwrap (),
                    cairo::Matrix::new (1.0, 2.0, 3.0, 4.0, 5.0, 6.0));

        assert_eq! (parse_transform ("matrix (1,2,3,4 5 6)").unwrap (),
                    cairo::Matrix::new (1.0, 2.0, 3.0, 4.0, 5.0, 6.0));

        assert_eq! (parse_transform ("matrix (1,2.25,-3.25e2,4 5 6)").unwrap (),
                    cairo::Matrix::new (1.0, 2.25, -325.0, 4.0, 5.0, 6.0));
    }

    #[test]
    #[ignore]
    fn parses_translate () {
        assert_eq! (parse_transform ("translate(-1 -2)").unwrap (),
                    cairo::Matrix::new (1.0, 0.0, 0.0, 1.0, -1.0, -2.0));

        assert_eq! (parse_transform ("translate(-1, -2)").unwrap (),
                    cairo::Matrix::new (1.0, 0.0, 0.0, 1.0, -1.0, -2.0));

        assert_eq! (parse_transform ("translate(-1)").unwrap (),
                    cairo::Matrix::new (1.0, 0.0, 0.0, 1.0, -1.0, 0.0));
    }

    #[test]
    #[ignore]
    fn parses_scale () {
        assert_eq! (parse_transform ("scale(-1 -2)").unwrap (),
                    cairo::Matrix::new (-1.0, 0.0, 0.0, -2.0, 0.0, 0.0));

        assert_eq! (parse_transform ("scale(-1, -2)").unwrap (),
                    cairo::Matrix::new (-1.0, 0.0, 0.0, -2.0, 0.0, 0.0));

        assert_eq! (parse_transform ("scale(-1)").unwrap (),
                    cairo::Matrix::new (-1.0, 0.0, 0.0, -1.0, 0.0, 0.0));
    }

    #[test]
    #[ignore]
    fn parses_rotate () {
        assert_eq! (parse_transform ("rotate (30)").unwrap (), make_rotation_matrix (30.0, 0.0, 0.0));
        assert_eq! (parse_transform ("rotate (30,-1,-2)").unwrap (), make_rotation_matrix (30.0, -1.0, -2.0));
        assert_eq! (parse_transform ("rotate (30, -1, -2)").unwrap (), make_rotation_matrix (30.0, -1.0, -2.0));
    }

    fn make_skew_x_matrix (angle_degrees: f64) -> cairo::Matrix {
        let a = angle_degrees * PI / 180.0;
        cairo::Matrix::new (1.0,      0.0,
                            a.tan (), 1.0,
                            0.0, 0.0)
    }

    fn make_skew_y_matrix (angle_degrees: f64) -> cairo::Matrix {
        let mut m = make_skew_x_matrix (angle_degrees);
        m.yx = m.xy;
        m.xy = 0.0;
        m
    }

    #[test]
    #[ignore]
    fn parses_skew_x () {
        assert_eq! (parse_transform ("skewX (30)").unwrap (), make_skew_x_matrix (30.0));
    }

    #[test]
    #[ignore]
    fn parses_skew_y () {
        assert_eq! (parse_transform ("skewY (30)").unwrap (), make_skew_y_matrix (30.0));
    }

    #[test]
    #[ignore]
    fn parses_transform_list () {
        let t = cairo::Matrix::new (1.0, 0.0, 0.0, 1.0, 20.0, 30.0);
        let s = cairo::Matrix::new (10.0, 0.0, 0.0, 10.0, 0.0, 0.0);
        let r = make_rotation_matrix (30.0, 10.0, 10.0);

        assert_eq! (parse_transform ("scale(10)rotate(30, 10, 10)").unwrap (),
                    cairo::Matrix::multiply (&r, &s));

        assert_eq! (parse_transform ("translate(20, 30), scale (10)").unwrap (),
                    cairo::Matrix::multiply (&s, &t));

        let a = cairo::Matrix::multiply (&s, &t);
        assert_eq! (parse_transform ("translate(20, 30), scale (10) rotate (30 10 10)").unwrap (),
                    cairo::Matrix::multiply (&r, &a));
    }
}
<commit_msg>Compose multiple transforms into a transform list<commit_after>use ::cairo;
use ::glib::translate::*;
use ::glib_sys;
use ::libc;

#[cfg(test)]
use std::f64::consts::*;

use cairo::MatrixTrait;
use cssparser::{self, Parser, ParserInput};

use error::*;
use parsers::{ParseError, Parse, optional_comma};

impl Parse for cairo::Matrix {
    type Data = ();
    type Err = AttributeError;

    fn parse (s: &str, _: ()) -> Result<cairo::Matrix, AttributeError> {
        parse_transform (s)
    }
}

pub fn parse_transform(s: &str) -> Result<cairo::Matrix, AttributeError> {
    let matrix = parse_transform_list(s)?;

    matrix.try_invert ().map (|_| matrix)
        .map_err (|_| AttributeError::Value ("invalid transformation matrix".to_string ()))
}

fn parse_transform_list(s: &str) -> Result<cairo::Matrix, AttributeError> {
    let mut input = ParserInput::new(s);
    let mut parser = Parser::new(&mut input);

    let mut matrix = cairo::Matrix::identity();

    loop {
        let m = parse_transform_command(&mut parser)?;
        matrix = cairo::Matrix::multiply(&m, &matrix);

        if parser.is_exhausted() {
            break;
        }

        optional_comma(&mut parser);
    }

    Ok(matrix)
}

fn parse_transform_command(parser: &mut Parser) -> Result<cairo::Matrix, AttributeError> {
    let xform = parser.expect_ident_cloned()?;
    let _ = parser.expect_parenthesis_block()?;

    match xform.as_ref() {
        "matrix"    => parse_matrix_args(parser),
        "translate" => parse_translate_args(parser),
        "scale"     => parse_scale_args(parser),
        "rotate"    => parse_rotate_args(parser),
        "skewX"     => parse_skewX_args(parser),
        "skewY"     => parse_skewY_args(parser),
        _           => Err(AttributeError::from(ParseError::new("expected matrix|translate|scale|rotate|skewX|skewY"))),
    }
}

fn parse_matrix_args(parser: &mut Parser) -> Result<cairo::Matrix, AttributeError> {
    parser.parse_nested_block(|p| {
        let xx = p.expect_number()? as f64;
        optional_comma(p);

        let yx = p.expect_number()? as f64;
        optional_comma(p);

        let xy = p.expect_number()? as f64;
        optional_comma(p);

        let yy = p.expect_number()? as f64;
        optional_comma(p);

        let x0 = p.expect_number()? as f64;
        optional_comma(p);

        let y0 = p.expect_number()? as f64;

        Ok(cairo::Matrix::new(xx, yx, xy, yy, x0, y0))
    }).map_err(cssparser::ParseError::<()>::basic)
        .map_err(|e| AttributeError::from(e))
}

fn parse_translate_args(parser: &mut Parser) -> Result<cairo::Matrix, AttributeError> {
    unimplemented!();
}

fn parse_scale_args(parser: &mut Parser) -> Result<cairo::Matrix, AttributeError> {
    unimplemented!();
}

fn parse_rotate_args(parser: &mut Parser) -> Result<cairo::Matrix, AttributeError> {
    unimplemented!();
}

fn parse_skewX_args(parser: &mut Parser) -> Result<cairo::Matrix, AttributeError> {
    unimplemented!();
}

fn parse_skewY_args(parser: &mut Parser) -> Result<cairo::Matrix, AttributeError> {
    unimplemented!();
}

/*
pub fn parse_transform (s: &str) -> Result <cairo::Matrix, AttributeError> {
    let r = parse_TransformList (s);

    match r {
        Ok (m) => {
            m.try_invert ().map (|_| m)
                .map_err (|_| AttributeError::Value ("invalid transformation matrix".to_string ()))
        },

        Err (e) => {
            Err (AttributeError::Parse (ParseError::new (format! ("{:?}", e))))
        }
    }
}
*/

#[cfg(test)]
fn make_rotation_matrix (angle_degrees: f64, tx: f64, ty: f64) -> cairo::Matrix {
    let angle = angle_degrees * PI / 180.0;

    let mut m = cairo::Matrix::new (1.0, 0.0, 0.0, 1.0, tx, ty);

    let mut r = cairo::Matrix::identity ();
    r.rotate (angle);
    m = cairo::Matrix::multiply (&r, &m);

    m = cairo::Matrix::multiply (&cairo::Matrix::new (1.0, 0.0, 0.0, 1.0, -tx, -ty), &m);
    m
}

#[no_mangle]
pub extern fn rsvg_parse_transform (out_matrix: *mut cairo::Matrix, s: *const libc::c_char) -> glib_sys::gboolean {
    assert! (!out_matrix.is_null ());
    assert! (!s.is_null ());

    let string = unsafe { String::from_glib_none (s) };
    let matrix: &mut cairo::Matrix = unsafe { &mut *out_matrix };

    match parse_transform (&string) {
        Ok (m) => {
            *matrix = m;
            true.to_glib ()
        },

        Err (_) => {
            *matrix = cairo::Matrix::identity ();
            false.to_glib ()
        }
    }
}

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

    #[test]
    #[ignore]
    fn parses_valid_transform () {
        let t = cairo::Matrix::new (1.0, 0.0, 0.0, 1.0, 20.0, 30.0);
        let s = cairo::Matrix::new (10.0, 0.0, 0.0, 10.0, 0.0, 0.0);
        let r = make_rotation_matrix (30.0, 10.0, 10.0);

        let a = cairo::Matrix::multiply (&s, &t);
        assert_eq! (parse_transform ("translate(20, 30), scale (10) rotate (30 10 10)").unwrap (),
                    cairo::Matrix::multiply (&r, &a));
    }

    #[test]
    #[ignore]
    fn syntax_error_yields_parse_error () {
        match parse_transform ("foo") {
            Err (AttributeError::Parse (_)) => {},
            _ => { panic! (); }
        }

        match parse_transform ("matrix (1 2 3 4 5)") {
            Err (AttributeError::Parse (_)) => {},
            _ => { panic! (); }
        }
    }

    #[test]
    #[ignore]
    fn invalid_transform_yields_value_error () {
        match parse_transform ("matrix (0 0 0 0 0 0)") {
            Err (AttributeError::Value (_)) => {},
            _ => { panic! (); }
        }

        match parse_transform ("scale (0), translate (10, 10)") {
            Err (AttributeError::Value (_)) => {},
            _ => { panic! (); }
        }

        match parse_transform ("scale (0), skewX (90)") {
            Err (AttributeError::Value (_)) => {},
            _ => { panic! (); }
        }
    }
}

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

    #[test]
    fn parses_matrix () {
        assert_eq! (parse_transform ("matrix (1 2 3 4 5 6)").unwrap (),
                    cairo::Matrix::new (1.0, 2.0, 3.0, 4.0, 5.0, 6.0));

        assert_eq! (parse_transform ("matrix (1,2,3,4 5 6)").unwrap (),
                    cairo::Matrix::new (1.0, 2.0, 3.0, 4.0, 5.0, 6.0));

        assert_eq! (parse_transform ("matrix (1,2.25,-3.25e2,4 5 6)").unwrap (),
                    cairo::Matrix::new (1.0, 2.25, -325.0, 4.0, 5.0, 6.0));
    }

    #[test]
    #[ignore]
    fn parses_translate () {
        assert_eq! (parse_transform ("translate(-1 -2)").unwrap (),
                    cairo::Matrix::new (1.0, 0.0, 0.0, 1.0, -1.0, -2.0));

        assert_eq! (parse_transform ("translate(-1, -2)").unwrap (),
                    cairo::Matrix::new (1.0, 0.0, 0.0, 1.0, -1.0, -2.0));

        assert_eq! (parse_transform ("translate(-1)").unwrap (),
                    cairo::Matrix::new (1.0, 0.0, 0.0, 1.0, -1.0, 0.0));
    }

    #[test]
    #[ignore]
    fn parses_scale () {
        assert_eq! (parse_transform ("scale(-1 -2)").unwrap (),
                    cairo::Matrix::new (-1.0, 0.0, 0.0, -2.0, 0.0, 0.0));

        assert_eq! (parse_transform ("scale(-1, -2)").unwrap (),
                    cairo::Matrix::new (-1.0, 0.0, 0.0, -2.0, 0.0, 0.0));

        assert_eq! (parse_transform ("scale(-1)").unwrap (),
                    cairo::Matrix::new (-1.0, 0.0, 0.0, -1.0, 0.0, 0.0));
    }

    #[test]
    #[ignore]
    fn parses_rotate () {
        assert_eq! (parse_transform ("rotate (30)").unwrap (), make_rotation_matrix (30.0, 0.0, 0.0));
        assert_eq! (parse_transform ("rotate (30,-1,-2)").unwrap (), make_rotation_matrix (30.0, -1.0, -2.0));
        assert_eq! (parse_transform ("rotate (30, -1, -2)").unwrap (), make_rotation_matrix (30.0, -1.0, -2.0));
    }

    fn make_skew_x_matrix (angle_degrees: f64) -> cairo::Matrix {
        let a = angle_degrees * PI / 180.0;
        cairo::Matrix::new (1.0,      0.0,
                            a.tan (), 1.0,
                            0.0, 0.0)
    }

    fn make_skew_y_matrix (angle_degrees: f64) -> cairo::Matrix {
        let mut m = make_skew_x_matrix (angle_degrees);
        m.yx = m.xy;
        m.xy = 0.0;
        m
    }

    #[test]
    #[ignore]
    fn parses_skew_x () {
        assert_eq! (parse_transform ("skewX (30)").unwrap (), make_skew_x_matrix (30.0));
    }

    #[test]
    #[ignore]
    fn parses_skew_y () {
        assert_eq! (parse_transform ("skewY (30)").unwrap (), make_skew_y_matrix (30.0));
    }

    #[test]
    #[ignore]
    fn parses_transform_list () {
        let t = cairo::Matrix::new (1.0, 0.0, 0.0, 1.0, 20.0, 30.0);
        let s = cairo::Matrix::new (10.0, 0.0, 0.0, 10.0, 0.0, 0.0);
        let r = make_rotation_matrix (30.0, 10.0, 10.0);

        assert_eq! (parse_transform ("scale(10)rotate(30, 10, 10)").unwrap (),
                    cairo::Matrix::multiply (&r, &s));

        assert_eq! (parse_transform ("translate(20, 30), scale (10)").unwrap (),
                    cairo::Matrix::multiply (&s, &t));

        let a = cairo::Matrix::multiply (&s, &t);
        assert_eq! (parse_transform ("translate(20, 30), scale (10) rotate (30 10 10)").unwrap (),
                    cairo::Matrix::multiply (&r, &a));
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::env::{current_dir, vars};
use std::io::{self, Write};

use anyhow::{bail, Context, Result};
use log::error;

use pueue::message::*;
use pueue::platform::socket::*;
use pueue::protocol::*;
use pueue::settings::Settings;

use crate::cli::{Opt, SubCommand};
use crate::commands::edit::*;
use crate::commands::local_follow::*;
use crate::commands::restart::*;
use crate::output::*;

/// This struct contains the base logic for the client.
/// The client is responsible for connecting to the daemon, sending instructions
/// and interpreting their responses.
///
/// Most commands are a simple ping-pong. However, some commands require a more complex
/// communication pattern, such as the `follow` command, which continuously streams the output of a task.
pub struct Client {
    opt: Opt,
    settings: Settings,
    socket: Socket,
}

impl Client {
    /// Connect to the daemon, authorize via secret and return a new initialized Client.
    pub async fn new(settings: Settings, opt: Opt) -> Result<Self> {
        // // Commandline argument overwrites the configuration files values for address
        // let address = if let Some(address) = opt.address.clone() {
        //     address
        // } else {
        //     settings.client.daemon_address
        // };

        // Commandline argument overwrites the configuration files values for port
        let (unix_socket_path, port) = {
            // Return the unix socket path, if we're supposed to use it.
            if settings.shared.use_unix_socket {
                (Some(settings.shared.unix_socket_path.clone()), None)
            } else {
                // Otherwise use tcp sockets and a given port
                // Commandline argument overwrites the configuration files values for port.
                let port = if let Some(port) = opt.port.clone() {
                    port
                } else {
                    settings.shared.port.clone()
                };

                (None, Some(port))
            }
        };

        let mut socket = get_client(unix_socket_path, port).await?;

        // Send the secret to the daemon
        let secret = settings.shared.secret.clone().into_bytes();
        send_bytes(secret, &mut socket).await?;

        Ok(Client {
            opt,
            settings,
            socket,
        })
    }

    /// This is the function where the actual communication and logic starts.
    /// At this point everything is initialized, the connection is up and
    /// we can finally start doing stuff.
    ///
    /// The command handling is splitted into "simple" and "complex" commands.
    pub async fn start(&mut self) -> Result<()> {
        // Return early, if the command has already been handled.
        if self.handle_complex_command().await? {
            return Ok(());
        }

        // The handling of "generic" commands is encapsulated in this function.
        self.handle_simple_command().await?;

        Ok(())
    }

    /// Handle all complex client-side functionalities.
    /// Complex functionalities need some special handling and are contained
    /// in their own functions with their own communication code.
    /// Such functionalities includes reading local files, data streaming
    /// and sending multiple messages.
    ///
    /// Returns `true`, if the current command has been handled by this function.
    /// This indicates that the client can now shut down.
    async fn handle_complex_command(&mut self) -> Result<bool> {
        // This match handles all "complex" commands.
        match &self.opt.cmd {
            SubCommand::Edit { task_id, path } => {
                let message = edit(&mut self.socket, *task_id, *path).await?;
                self.handle_response(message);
                Ok(true)
            }
            SubCommand::Restart {
                task_ids,
                start_immediately,
                stashed,
                edit,
                path,
            } => {
                restart(
                    &mut self.socket,
                    task_ids.clone(),
                    *start_immediately,
                    *stashed,
                    *edit,
                    *path,
                )
                .await?;
                Ok(true)
            }
            SubCommand::Follow { task_id, err } => {
                // Simple log output follows for local logs don't need any communication with the daemon.
                // Thereby we handle this separately over here.
                if self.settings.client.read_local_logs {
                    local_follow(
                        &mut self.socket,
                        self.settings.shared.pueue_directory.clone(),
                        task_id,
                        *err,
                    )
                    .await?;
                    return Ok(true);
                }
                Ok(false)
            }
            _ => Ok(false),
        }
    }

    /// Handle logic that's super generic on the client-side.
    /// This always follows a singular ping-pong pattern.
    /// One message to the daemon, one response, Done.
    async fn handle_simple_command(&mut self) -> Result<()> {
        // Create the message that should be sent to the daemon
        // depending on the given commandline options.
        let message = self.get_message_from_opt()?;

        // Create the message payload and send it to the daemon.
        send_message(message, &mut self.socket).await?;

        // Check if we can receive the response from the daemon
        let mut response = receive_message(&mut self.socket).await?;

        // Check if we can receive the response from the daemon
        while self.handle_response(response) {
            response = receive_message(&mut self.socket).await?;
        }

        Ok(())
    }

    /// Most returned messages can be handled in a generic fashion.
    /// However, some commands require continuous receiving of messages (streaming).
    ///
    /// If this function returns `Ok(true)`, the parent function will continue to receive
    /// and handle messages from the daemon. Otherwise the client will simply exit.
    fn handle_response(&self, message: Message) -> bool {
        match message {
            Message::Success(text) => print_success(&text),
            Message::Failure(text) => print_error(&text),
            Message::StatusResponse(state) => print_state(state, &self.opt.cmd),
            Message::LogResponse(task_logs) => print_logs(task_logs, &self.opt.cmd, &self.settings),
            Message::Stream(text) => {
                print!("{}", text);
                io::stdout().flush().unwrap();
                return true;
            }
            _ => error!("Received unhandled response message"),
        };

        false
    }

    /// Prints a warning and prompt for given action and tasks.
    /// Returns `Ok(())` if the action was confirmed.
    fn handle_user_confirmation(&self, action: &str, task_ids: &Vec<usize>) -> Result<()> {
        // printing warning and prompt
        println!(
            "You are trying to {}: {}",
            action,
            task_ids
                .into_iter()
                .map(|t| format!("task{}", t.to_string()))
                .collect::<Vec<String>>()
                .join(", ")
        );
        print!("\nDo you want to continue [Y/n]: ");

        let mut input = String::new();
        io::stdout().flush().unwrap();
        io::stdin().read_line(&mut input)?;

        if input.starts_with("n") || input.starts_with("N") {
            bail!("User did not confirm");
        }

        Ok(())
    }

    /// Convert the cli command into the message that's being sent to the server,
    /// so it can be understood by the daemon.
    fn get_message_from_opt(&self) -> Result<Message> {
        match &self.opt.cmd {
            SubCommand::Add {
                command,
                start_immediately,
                stashed,
                group,
                delay_until,
                dependencies,
            } => {
                let cwd_pathbuf = current_dir()?;
                let cwd = cwd_pathbuf
                    .to_str()
                    .context("Cannot parse current working directory (Invalid utf8?)")?;

                let mut envs = HashMap::new();
                // Save all environment variables for later injection into the started task
                for (key, value) in vars() {
                    envs.insert(key, value);
                }

                Ok(Message::Add(AddMessage {
                    command: command.join(" "),
                    path: cwd.to_string(),
                    envs,
                    start_immediately: *start_immediately,
                    stashed: *stashed,
                    group: group.clone(),
                    enqueue_at: *delay_until,
                    dependencies: dependencies.to_vec(),
                    ignore_aliases: false,
                }))
            }
            SubCommand::Remove { task_ids } => {
                if self.settings.client.print_remove_warnings {
                    match self.handle_user_confirmation("remove", task_ids) {
                        Ok(_) => (),
                        Err(_) => std::process::exit(1),
                    }
                }
                Ok(Message::Remove(task_ids.clone()))
            }
            SubCommand::Stash { task_ids } => Ok(Message::Stash(task_ids.clone())),
            SubCommand::Switch {
                task_id_1,
                task_id_2,
            } => {
                let message = SwitchMessage {
                    task_id_1: *task_id_1,
                    task_id_2: *task_id_2,
                };
                Ok(Message::Switch(message))
            }
            SubCommand::Enqueue {
                task_ids,
                delay_until,
            } => {
                let message = EnqueueMessage {
                    task_ids: task_ids.clone(),
                    enqueue_at: *delay_until,
                };
                Ok(Message::Enqueue(message))
            }
            SubCommand::Start {
                task_ids,
                group,
                all,
                children,
            } => {
                let message = StartMessage {
                    task_ids: task_ids.clone(),
                    group: group.clone(),
                    all: *all,
                    children: *children,
                };
                Ok(Message::Start(message))
            }
            SubCommand::Pause {
                task_ids,
                group,
                wait,
                all,
                children,
            } => {
                let message = PauseMessage {
                    task_ids: task_ids.clone(),
                    group: group.clone(),
                    wait: *wait,
                    all: *all,
                    children: *children,
                };
                Ok(Message::Pause(message))
            }
            SubCommand::Kill {
                task_ids,
                group,
                default,
                all,
                children,
            } => {
                if self.settings.client.print_remove_warnings {
                    match self.handle_user_confirmation("kill", task_ids) {
                        Ok(_) => (),
                        Err(_) => std::process::exit(1),
                    }
                }
                let message = KillMessage {
                    task_ids: task_ids.clone(),
                    group: group.clone(),
                    default: *default,
                    all: *all,
                    children: *children,
                };
                Ok(Message::Kill(message))
            }
            SubCommand::Send { task_id, input } => {
                let message = SendMessage {
                    task_id: *task_id,
                    input: input.clone(),
                };
                Ok(Message::Send(message))
            }
            SubCommand::Group { add, remove } => {
                let message = GroupMessage {
                    add: add.clone(),
                    remove: remove.clone(),
                };
                Ok(Message::Group(message))
            }
            SubCommand::Status { .. } => Ok(Message::Status),
            SubCommand::Log { task_ids, .. } => {
                let message = LogRequestMessage {
                    task_ids: task_ids.clone(),
                    send_logs: !self.settings.client.read_local_logs,
                };
                Ok(Message::Log(message))
            }
            SubCommand::Follow { task_id, err } => {
                let message = StreamRequestMessage {
                    task_id: *task_id,
                    err: *err,
                };
                Ok(Message::StreamRequest(message))
            }
            SubCommand::Clean => Ok(Message::Clean),
            SubCommand::Reset { children } => Ok(Message::Reset(*children)),
            SubCommand::Shutdown => Ok(Message::DaemonShutdown),
            SubCommand::Parallel {
                parallel_tasks,
                group,
            } => {
                let message = ParallelMessage {
                    parallel_tasks: *parallel_tasks,
                    group: group.clone(),
                };
                Ok(Message::Parallel(message))
            }
            SubCommand::Completions { .. } => bail!("Completions have to be handled earlier"),
            SubCommand::Restart { .. } => bail!("Restarts have to be handled earlier"),
            SubCommand::Edit { .. } => bail!("Edits have to be handled earlier"),
        }
    }
}
<commit_msg>Fix clippy warnings<commit_after>use std::collections::HashMap;
use std::env::{current_dir, vars};
use std::io::{self, Write};

use anyhow::{bail, Context, Result};
use log::error;

use pueue::message::*;
use pueue::platform::socket::*;
use pueue::protocol::*;
use pueue::settings::Settings;

use crate::cli::{Opt, SubCommand};
use crate::commands::edit::*;
use crate::commands::local_follow::*;
use crate::commands::restart::*;
use crate::output::*;

/// This struct contains the base logic for the client.
/// The client is responsible for connecting to the daemon, sending instructions
/// and interpreting their responses.
///
/// Most commands are a simple ping-pong. However, some commands require a more complex
/// communication pattern, such as the `follow` command, which continuously streams the output of a task.
pub struct Client {
    opt: Opt,
    settings: Settings,
    socket: Socket,
}

impl Client {
    /// Connect to the daemon, authorize via secret and return a new initialized Client.
    pub async fn new(settings: Settings, opt: Opt) -> Result<Self> {
        // // Commandline argument overwrites the configuration files values for address
        // let address = if let Some(address) = opt.address.clone() {
        //     address
        // } else {
        //     settings.client.daemon_address
        // };

        // Commandline argument overwrites the configuration files values for port
        let (unix_socket_path, port) = {
            // Return the unix socket path, if we're supposed to use it.
            if settings.shared.use_unix_socket {
                (Some(settings.shared.unix_socket_path.clone()), None)
            } else {
                // Otherwise use tcp sockets and a given port
                // Commandline argument overwrites the configuration files values for port.
                let port = if let Some(port) = opt.port.clone() {
                    port
                } else {
                    settings.shared.port.clone()
                };

                (None, Some(port))
            }
        };

        let mut socket = get_client(unix_socket_path, port).await?;

        // Send the secret to the daemon
        let secret = settings.shared.secret.clone().into_bytes();
        send_bytes(secret, &mut socket).await?;

        Ok(Client {
            opt,
            settings,
            socket,
        })
    }

    /// This is the function where the actual communication and logic starts.
    /// At this point everything is initialized, the connection is up and
    /// we can finally start doing stuff.
    ///
    /// The command handling is splitted into "simple" and "complex" commands.
    pub async fn start(&mut self) -> Result<()> {
        // Return early, if the command has already been handled.
        if self.handle_complex_command().await? {
            return Ok(());
        }

        // The handling of "generic" commands is encapsulated in this function.
        self.handle_simple_command().await?;

        Ok(())
    }

    /// Handle all complex client-side functionalities.
    /// Complex functionalities need some special handling and are contained
    /// in their own functions with their own communication code.
    /// Such functionalities includes reading local files, data streaming
    /// and sending multiple messages.
    ///
    /// Returns `true`, if the current command has been handled by this function.
    /// This indicates that the client can now shut down.
    async fn handle_complex_command(&mut self) -> Result<bool> {
        // This match handles all "complex" commands.
        match &self.opt.cmd {
            SubCommand::Edit { task_id, path } => {
                let message = edit(&mut self.socket, *task_id, *path).await?;
                self.handle_response(message);
                Ok(true)
            }
            SubCommand::Restart {
                task_ids,
                start_immediately,
                stashed,
                edit,
                path,
            } => {
                restart(
                    &mut self.socket,
                    task_ids.clone(),
                    *start_immediately,
                    *stashed,
                    *edit,
                    *path,
                )
                .await?;
                Ok(true)
            }
            SubCommand::Follow { task_id, err } => {
                // Simple log output follows for local logs don't need any communication with the daemon.
                // Thereby we handle this separately over here.
                if self.settings.client.read_local_logs {
                    local_follow(
                        &mut self.socket,
                        self.settings.shared.pueue_directory.clone(),
                        task_id,
                        *err,
                    )
                    .await?;
                    return Ok(true);
                }
                Ok(false)
            }
            _ => Ok(false),
        }
    }

    /// Handle logic that's super generic on the client-side.
    /// This always follows a singular ping-pong pattern.
    /// One message to the daemon, one response, Done.
    async fn handle_simple_command(&mut self) -> Result<()> {
        // Create the message that should be sent to the daemon
        // depending on the given commandline options.
        let message = self.get_message_from_opt()?;

        // Create the message payload and send it to the daemon.
        send_message(message, &mut self.socket).await?;

        // Check if we can receive the response from the daemon
        let mut response = receive_message(&mut self.socket).await?;

        // Check if we can receive the response from the daemon
        while self.handle_response(response) {
            response = receive_message(&mut self.socket).await?;
        }

        Ok(())
    }

    /// Most returned messages can be handled in a generic fashion.
    /// However, some commands require continuous receiving of messages (streaming).
    ///
    /// If this function returns `Ok(true)`, the parent function will continue to receive
    /// and handle messages from the daemon. Otherwise the client will simply exit.
    fn handle_response(&self, message: Message) -> bool {
        match message {
            Message::Success(text) => print_success(&text),
            Message::Failure(text) => print_error(&text),
            Message::StatusResponse(state) => print_state(state, &self.opt.cmd),
            Message::LogResponse(task_logs) => print_logs(task_logs, &self.opt.cmd, &self.settings),
            Message::Stream(text) => {
                print!("{}", text);
                io::stdout().flush().unwrap();
                return true;
            }
            _ => error!("Received unhandled response message"),
        };

        false
    }

    /// Prints a warning and prompt for given action and tasks.
    /// Returns `Ok(())` if the action was confirmed.
    fn handle_user_confirmation(&self, action: &str, task_ids: &[usize]) -> Result<()> {
        // printing warning and prompt
        println!(
            "You are trying to {}: {}",
            action,
            task_ids
                .iter()
                .map(|t| format!("task{}", t.to_string()))
                .collect::<Vec<String>>()
                .join(", ")
        );
        print!("\nDo you want to continue [Y/n]: ");

        let mut input = String::new();
        io::stdout().flush().unwrap();
        io::stdin().read_line(&mut input)?;

        if input.starts_with('n') || input.starts_with('N') {
            bail!("User did not confirm");
        }

        Ok(())
    }

    /// Convert the cli command into the message that's being sent to the server,
    /// so it can be understood by the daemon.
    fn get_message_from_opt(&self) -> Result<Message> {
        match &self.opt.cmd {
            SubCommand::Add {
                command,
                start_immediately,
                stashed,
                group,
                delay_until,
                dependencies,
            } => {
                let cwd_pathbuf = current_dir()?;
                let cwd = cwd_pathbuf
                    .to_str()
                    .context("Cannot parse current working directory (Invalid utf8?)")?;

                let mut envs = HashMap::new();
                // Save all environment variables for later injection into the started task
                for (key, value) in vars() {
                    envs.insert(key, value);
                }

                Ok(Message::Add(AddMessage {
                    command: command.join(" "),
                    path: cwd.to_string(),
                    envs,
                    start_immediately: *start_immediately,
                    stashed: *stashed,
                    group: group.clone(),
                    enqueue_at: *delay_until,
                    dependencies: dependencies.to_vec(),
                    ignore_aliases: false,
                }))
            }
            SubCommand::Remove { task_ids } => {
                if self.settings.client.print_remove_warnings {
                    match self.handle_user_confirmation("remove", task_ids) {
                        Ok(_) => (),
                        Err(_) => std::process::exit(1),
                    }
                }
                Ok(Message::Remove(task_ids.clone()))
            }
            SubCommand::Stash { task_ids } => Ok(Message::Stash(task_ids.clone())),
            SubCommand::Switch {
                task_id_1,
                task_id_2,
            } => {
                let message = SwitchMessage {
                    task_id_1: *task_id_1,
                    task_id_2: *task_id_2,
                };
                Ok(Message::Switch(message))
            }
            SubCommand::Enqueue {
                task_ids,
                delay_until,
            } => {
                let message = EnqueueMessage {
                    task_ids: task_ids.clone(),
                    enqueue_at: *delay_until,
                };
                Ok(Message::Enqueue(message))
            }
            SubCommand::Start {
                task_ids,
                group,
                all,
                children,
            } => {
                let message = StartMessage {
                    task_ids: task_ids.clone(),
                    group: group.clone(),
                    all: *all,
                    children: *children,
                };
                Ok(Message::Start(message))
            }
            SubCommand::Pause {
                task_ids,
                group,
                wait,
                all,
                children,
            } => {
                let message = PauseMessage {
                    task_ids: task_ids.clone(),
                    group: group.clone(),
                    wait: *wait,
                    all: *all,
                    children: *children,
                };
                Ok(Message::Pause(message))
            }
            SubCommand::Kill {
                task_ids,
                group,
                default,
                all,
                children,
            } => {
                if self.settings.client.print_remove_warnings {
                    match self.handle_user_confirmation("kill", task_ids) {
                        Ok(_) => (),
                        Err(_) => std::process::exit(1),
                    }
                }
                let message = KillMessage {
                    task_ids: task_ids.clone(),
                    group: group.clone(),
                    default: *default,
                    all: *all,
                    children: *children,
                };
                Ok(Message::Kill(message))
            }
            SubCommand::Send { task_id, input } => {
                let message = SendMessage {
                    task_id: *task_id,
                    input: input.clone(),
                };
                Ok(Message::Send(message))
            }
            SubCommand::Group { add, remove } => {
                let message = GroupMessage {
                    add: add.clone(),
                    remove: remove.clone(),
                };
                Ok(Message::Group(message))
            }
            SubCommand::Status { .. } => Ok(Message::Status),
            SubCommand::Log { task_ids, .. } => {
                let message = LogRequestMessage {
                    task_ids: task_ids.clone(),
                    send_logs: !self.settings.client.read_local_logs,
                };
                Ok(Message::Log(message))
            }
            SubCommand::Follow { task_id, err } => {
                let message = StreamRequestMessage {
                    task_id: *task_id,
                    err: *err,
                };
                Ok(Message::StreamRequest(message))
            }
            SubCommand::Clean => Ok(Message::Clean),
            SubCommand::Reset { children } => Ok(Message::Reset(*children)),
            SubCommand::Shutdown => Ok(Message::DaemonShutdown),
            SubCommand::Parallel {
                parallel_tasks,
                group,
            } => {
                let message = ParallelMessage {
                    parallel_tasks: *parallel_tasks,
                    group: group.clone(),
                };
                Ok(Message::Parallel(message))
            }
            SubCommand::Completions { .. } => bail!("Completions have to be handled earlier"),
            SubCommand::Restart { .. } => bail!("Restarts have to be handled earlier"),
            SubCommand::Edit { .. } => bail!("Edits have to be handled earlier"),
        }
    }
}
<|endoftext|>


<commit_before>use std::net::{SocketAddr, SocketAddrV4, SocketAddrV6, IpAddr};
use std::fmt;

pub enum ProxyProtocolHeader {
  V1(HeaderV1),
  V2(HeaderV2),
}

impl ProxyProtocolHeader {
  // Use this method to writte the header in the backend socket
  pub fn into_bytes(&self) -> Vec<u8> {
    match *self {
      ProxyProtocolHeader::V1(ref header) => header.into_bytes(),
      ProxyProtocolHeader::V2(ref header) => header.into_bytes(),
    }
  }
}

/// Indicate the proxied INET protocol and family
#[derive(Debug, PartialEq, Eq)]
pub enum ProtocolSupportedV1 {
  TCP4,    // TCP over IPv4
  TCP6,    // TCP over IPv6
  UNKNOWN, // unsupported,or unknown protocols
}

impl fmt::Display for ProtocolSupportedV1 {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    match *self {
      ProtocolSupportedV1::TCP4 => write!(f, "TCP4"),
      ProtocolSupportedV1::TCP6 => write!(f, "TCP6"),
      ProtocolSupportedV1::UNKNOWN => write!(f, "UNKNOWN"),
    }
  }
}

/// Proxy Protocol header for version 1 (text version)
/// Example:
/// - TCP/IPv4: `PROXY TCP4 255.255.255.255 255.255.255.255 65535 65535\r\n`
/// - TCP/IPv6: `PROXY TCP6 ffff:f...f:ffff ffff:f...f:ffff 65535 65535\r\n`
/// - Unknown: `PROXY UNKNOWN\r\n`
#[derive(Debug)]
pub struct HeaderV1 {
  protocol: ProtocolSupportedV1,
  addr_src: SocketAddr,
  addr_dst: SocketAddr,
}

const PROXY_PROTO_IDENTIFIER: &'static str = "PROXY";

impl HeaderV1 {

  pub fn new (addr_src: SocketAddr, addr_dst: SocketAddr) -> Self {
    let protocol = if addr_dst.is_ipv6() {
      ProtocolSupportedV1::TCP6
    } else if addr_dst.is_ipv4() {
      ProtocolSupportedV1::TCP4
    } else {
      ProtocolSupportedV1::UNKNOWN
    };

    HeaderV1 {
      protocol,
      addr_src,
      addr_dst,
    }
  }

  pub fn into_bytes(&self) -> Vec<u8> {
    if self.protocol.eq(&ProtocolSupportedV1::UNKNOWN) {
      format!("{} {}\r\n",
          PROXY_PROTO_IDENTIFIER,
          self.protocol,
      ).into_bytes()
    } else {
      format!("{} {} {} {} {} {}\r\n",
        PROXY_PROTO_IDENTIFIER,
        self.protocol,
        self.addr_src.ip(),
        self.addr_dst.ip(),
        self.addr_src.port(),
        self.addr_dst.port(),
      ).into_bytes()
    }
  }
}

#[cfg(test)]
mod test {

  use super::*;
  use std::net::{IpAddr, Ipv4Addr, SocketAddr, Ipv6Addr};

  #[test]
  fn it_should_return_a_correct_header_with_ipv4() {
    let header = HeaderV1::new(
      SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 80),
      SocketAddr::new(IpAddr::V4(Ipv4Addr::new(172, 17, 40, 59)), 80)
    );

    let header_to_cmp = "PROXY TCP4 127.0.0.1 172.17.40.59 80 80\r\n".as_bytes();

    assert_eq!(header_to_cmp, &header.into_bytes()[..]);
  }

  #[test]
  fn it_should_return_a_correct_header_with_ipv6() {
    let header = HeaderV1::new(
      SocketAddr::new(IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 0xffff)), 80),
      SocketAddr::new(IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0x9c, 0x76)), 80)
    );

    let header_to_cmp = "PROXY TCP6 ::0.0.255.255 ::0.156.0.118 80 80\r\n".as_bytes();
    assert_eq!(header_to_cmp, &header.into_bytes()[..]);
  }

  // #[test]
  // fn it_should_return_the_an_unknown_header() {
  //   let header = HeaderV1::new(ProtocolSupportedV1::UNKNOWN, None, None);

  //   let header_to_cmp = "PROXY UNKNOWN\r\n".as_bytes();

  //   assert_eq!(header_to_cmp, &header.into_bytes()[..]);
  // }
}

pub struct HeaderV2 {
  signature: [u8; 12], // hex 0D 0A 0D 0A 00 0D 0A 51 55 49 54 0A
  ver_and_cmd: u8,     // protocol version and command
  family: u8,          // protocol family and address
  len: u16,            // number of following bytes part of the header
  addr: ProxyAddr,
}

impl HeaderV2 {
  pub fn new(addr_src: SocketAddr, addr_dst: SocketAddr) -> Self {
    let addr = ProxyAddr::from(addr_src, addr_dst);

    HeaderV2 {
      signature: [0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A],
      ver_and_cmd: 0x20,
      family: get_family(&addr),
      len: addr.len(),
      addr,
    }
  }

  pub fn into_bytes(&self) -> Vec<u8> {
    let mut header = Vec::new();
    header.extend_from_slice(&self.signature);
    header.push(self.ver_and_cmd);
    header.push(self.family);
    header.extend_from_slice(&u16_to_array_of_u8(self.len));
    header.extend_from_slice(&self.addr.into_bytes());
    header
  }
}

pub enum ProxyAddr {
  Ipv4Addr {
    src_addr: SocketAddrV4,
    dst_addr: SocketAddrV4,
  },
  Ipv6Addr {
    src_addr: SocketAddrV6,
    dst_addr: SocketAddrV6,
  },
  UnixAddr {
    src_addr: [u8; 108],
    dst_addr: [u8; 108],
  },
  AfUnspec,
}

impl ProxyAddr {

  pub fn from(addr_src: SocketAddr, addr_dst: SocketAddr) -> Self {
    match (addr_src, addr_dst) {
      (SocketAddr::V4(addr_ipv4_src), SocketAddr::V4(addr_ipv4_dst)) => {
        ProxyAddr::Ipv4Addr{
          src_addr: addr_ipv4_src,
          dst_addr: addr_ipv4_dst,
        }
      },
      (SocketAddr::V6(addr_ipv6_src), SocketAddr::V6(addr_ipv6_dst)) => {
        ProxyAddr::Ipv6Addr{
          src_addr: addr_ipv6_src,
          dst_addr: addr_ipv6_dst,
        }
      },
      _ => ProxyAddr::AfUnspec,
    }
  }

  fn len(&self) -> u16 {
    match *self {
      ProxyAddr::Ipv4Addr{ src_addr: _, dst_addr: _ } => 12,
      ProxyAddr::Ipv6Addr{ src_addr: _, dst_addr: _ } => 36,
      ProxyAddr::UnixAddr{ src_addr: _, dst_addr: _ } => 216,
      ProxyAddr::AfUnspec => 0,
    }
  }

  fn into_bytes(&self) -> Vec<u8> {
    let mut res = Vec::new();

    match *self {
      ProxyAddr::Ipv4Addr{ src_addr, dst_addr } => {
        res.extend_from_slice(&src_addr.ip().octets());
        res.extend_from_slice(&dst_addr.ip().octets());
        res.extend_from_slice(&u16_to_array_of_u8(src_addr.port()));
        res.extend_from_slice(&u16_to_array_of_u8(dst_addr.port()));
      },
      ProxyAddr::Ipv6Addr{ src_addr, dst_addr } => {
        res.extend_from_slice(&src_addr.ip().octets());
        res.extend_from_slice(&dst_addr.ip().octets());
        res.extend_from_slice(&u16_to_array_of_u8(src_addr.port()));
        res.extend_from_slice(&u16_to_array_of_u8(dst_addr.port()));
      },
      ProxyAddr::UnixAddr{ src_addr, dst_addr } =>{
        res.extend_from_slice(&src_addr);
        res.extend_from_slice(&dst_addr);
      },
      ProxyAddr::AfUnspec => {},
    };

    res
  }
}

fn get_family(addr: &ProxyAddr) -> u8 {
  match addr {
    &ProxyAddr::Ipv4Addr{ src_addr: _, dst_addr: _ } => 0x10 | 0x01, // AF_INET  = 1 + STREAM = 1
    &ProxyAddr::Ipv6Addr{ src_addr: _, dst_addr: _ } => 0x20 | 0x01, // AF_INET6 = 2 + STREAM = 1
    &ProxyAddr::UnixAddr{ src_addr: _, dst_addr: _ } => 0x30 | 0x01, // AF_UNIX  = 3 + STREAM = 1
    &ProxyAddr::AfUnspec => 0x00, // AF_UNSPEC + UNSPEC
  }
}

fn u16_to_array_of_u8(x: u16) -> [u8; 2] {
    let b1 : u8 = ((x >> 8) & 0xff) as u8;
    let b2 : u8 = (x & 0xff) as u8;
    return [b1, b2]
}

#[cfg(test)]
mod testV2 {

  use super::*;
  use std::net::{Ipv4Addr, Ipv6Addr};

  #[test]
  fn test_u16_to_array_of_u8() {
    let val_u16: u16 = 65534;
    let expected = [0xff, 0xfe];
    assert_eq!(expected, u16_to_array_of_u8(val_u16));
  }

  #[test]
  fn test_deserialize_tcp_ipv4_proxy_protocol_header() {
    let src_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(125, 25, 10, 1)), 8080);
    let dst_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 4, 5, 8)), 4200);

    let header = HeaderV2::new(src_addr, dst_addr);
    let expected = &[
      0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A, // MAGIC header
      0x20,                                                                   // Version 2 and command LOCAL
      0x11,                                                                   // family AF_UNIX with IPv4
      0x00, 0x0C,                                                             // address sizes = 12
      0x7D, 0x19, 0x0A, 0x01,                                                 // source address
      0x0A, 0x04, 0x05, 0x08,                                                 // destination address
      0x1F, 0x90,                                                             // source port
      0x10, 0x68,                                                             // destination port
    ];

    assert_eq!(expected, &header.into_bytes()[..]);
  }

  #[test]
  fn test_deserialize_tcp_ipv6_proxy_protocol_header() {
    let src_addr = SocketAddr::new(IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1)), 8080);
    let dst_addr = SocketAddr::new(IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1)), 4200);

    let header = HeaderV2::new(src_addr, dst_addr);
    let expected = vec![
      0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A,                         // MAGIC header
      0x20,                                                                                // Version 2 and command LOCAL
      0x21,                                                                                // family AF_UNIX with IPv6
      0x00, 0x24,                                                                                // address sizes = 36
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, // source address
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, // destination address
      0x1F, 0x90,                                                                                // source port
      0x10, 0x68,                                                                                // destination port
    ];

    assert_eq!(&expected[..], &header.into_bytes()[..]);
  }
}<commit_msg>Set the capacity of the Vec create on proxy protocol header.into_bytes<commit_after>use std::net::{SocketAddr, SocketAddrV4, SocketAddrV6, IpAddr};
use std::fmt;

pub enum ProxyProtocolHeader {
  V1(HeaderV1),
  V2(HeaderV2),
}

impl ProxyProtocolHeader {
  // Use this method to writte the header in the backend socket
  pub fn into_bytes(&self) -> Vec<u8> {
    match *self {
      ProxyProtocolHeader::V1(ref header) => header.into_bytes(),
      ProxyProtocolHeader::V2(ref header) => header.into_bytes(),
    }
  }
}

/// Indicate the proxied INET protocol and family
#[derive(Debug, PartialEq, Eq)]
pub enum ProtocolSupportedV1 {
  TCP4,    // TCP over IPv4
  TCP6,    // TCP over IPv6
  UNKNOWN, // unsupported,or unknown protocols
}

impl fmt::Display for ProtocolSupportedV1 {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    match *self {
      ProtocolSupportedV1::TCP4 => write!(f, "TCP4"),
      ProtocolSupportedV1::TCP6 => write!(f, "TCP6"),
      ProtocolSupportedV1::UNKNOWN => write!(f, "UNKNOWN"),
    }
  }
}

/// Proxy Protocol header for version 1 (text version)
/// Example:
/// - TCP/IPv4: `PROXY TCP4 255.255.255.255 255.255.255.255 65535 65535\r\n`
/// - TCP/IPv6: `PROXY TCP6 ffff:f...f:ffff ffff:f...f:ffff 65535 65535\r\n`
/// - Unknown: `PROXY UNKNOWN\r\n`
#[derive(Debug)]
pub struct HeaderV1 {
  protocol: ProtocolSupportedV1,
  addr_src: SocketAddr,
  addr_dst: SocketAddr,
}

const PROXY_PROTO_IDENTIFIER: &'static str = "PROXY";

impl HeaderV1 {

  pub fn new (addr_src: SocketAddr, addr_dst: SocketAddr) -> Self {
    let protocol = if addr_dst.is_ipv6() {
      ProtocolSupportedV1::TCP6
    } else if addr_dst.is_ipv4() {
      ProtocolSupportedV1::TCP4
    } else {
      ProtocolSupportedV1::UNKNOWN
    };

    HeaderV1 {
      protocol,
      addr_src,
      addr_dst,
    }
  }

  pub fn into_bytes(&self) -> Vec<u8> {
    if self.protocol.eq(&ProtocolSupportedV1::UNKNOWN) {
      format!("{} {}\r\n",
          PROXY_PROTO_IDENTIFIER,
          self.protocol,
      ).into_bytes()
    } else {
      format!("{} {} {} {} {} {}\r\n",
        PROXY_PROTO_IDENTIFIER,
        self.protocol,
        self.addr_src.ip(),
        self.addr_dst.ip(),
        self.addr_src.port(),
        self.addr_dst.port(),
      ).into_bytes()
    }
  }
}

#[cfg(test)]
mod test {

  use super::*;
  use std::net::{IpAddr, Ipv4Addr, SocketAddr, Ipv6Addr};

  #[test]
  fn it_should_return_a_correct_header_with_ipv4() {
    let header = HeaderV1::new(
      SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 80),
      SocketAddr::new(IpAddr::V4(Ipv4Addr::new(172, 17, 40, 59)), 80)
    );

    let header_to_cmp = "PROXY TCP4 127.0.0.1 172.17.40.59 80 80\r\n".as_bytes();

    assert_eq!(header_to_cmp, &header.into_bytes()[..]);
  }

  #[test]
  fn it_should_return_a_correct_header_with_ipv6() {
    let header = HeaderV1::new(
      SocketAddr::new(IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 0xffff)), 80),
      SocketAddr::new(IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0x9c, 0x76)), 80)
    );

    let header_to_cmp = "PROXY TCP6 ::0.0.255.255 ::0.156.0.118 80 80\r\n".as_bytes();
    assert_eq!(header_to_cmp, &header.into_bytes()[..]);
  }

  // #[test]
  // fn it_should_return_the_an_unknown_header() {
  //   let header = HeaderV1::new(ProtocolSupportedV1::UNKNOWN, None, None);

  //   let header_to_cmp = "PROXY UNKNOWN\r\n".as_bytes();

  //   assert_eq!(header_to_cmp, &header.into_bytes()[..]);
  // }
}

pub struct HeaderV2 {
  signature: [u8; 12], // hex 0D 0A 0D 0A 00 0D 0A 51 55 49 54 0A
  ver_and_cmd: u8,     // protocol version and command
  family: u8,          // protocol family and address
  len: u16,            // number of following bytes part of the header
  addr: ProxyAddr,
}

impl HeaderV2 {
  pub fn new(addr_src: SocketAddr, addr_dst: SocketAddr) -> Self {
    let addr = ProxyAddr::from(addr_src, addr_dst);

    HeaderV2 {
      signature: [0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A],
      ver_and_cmd: 0x20,
      family: get_family(&addr),
      len: addr.len(),
      addr,
    }
  }

  pub fn into_bytes(&self) -> Vec<u8> {
    let mut header = Vec::with_capacity(self.len());
    header.extend_from_slice(&self.signature);
    header.push(self.ver_and_cmd);
    header.push(self.family);
    header.extend_from_slice(&u16_to_array_of_u8(self.len));
    self.addr.into_bytes(&mut header);
    header
  }

  pub fn len(&self) -> usize {
    // signature + ver_and_cmd + family + len + addr
    self.signature.len() + 1 + 1 + 2 + self.addr.len() as usize
  }
}

pub enum ProxyAddr {
  Ipv4Addr {
    src_addr: SocketAddrV4,
    dst_addr: SocketAddrV4,
  },
  Ipv6Addr {
    src_addr: SocketAddrV6,
    dst_addr: SocketAddrV6,
  },
  UnixAddr {
    src_addr: [u8; 108],
    dst_addr: [u8; 108],
  },
  AfUnspec,
}

impl ProxyAddr {

  pub fn from(addr_src: SocketAddr, addr_dst: SocketAddr) -> Self {
    match (addr_src, addr_dst) {
      (SocketAddr::V4(addr_ipv4_src), SocketAddr::V4(addr_ipv4_dst)) => {
        ProxyAddr::Ipv4Addr{
          src_addr: addr_ipv4_src,
          dst_addr: addr_ipv4_dst,
        }
      },
      (SocketAddr::V6(addr_ipv6_src), SocketAddr::V6(addr_ipv6_dst)) => {
        ProxyAddr::Ipv6Addr{
          src_addr: addr_ipv6_src,
          dst_addr: addr_ipv6_dst,
        }
      },
      _ => ProxyAddr::AfUnspec,
    }
  }

  fn len(&self) -> u16 {
    match *self {
      ProxyAddr::Ipv4Addr{ src_addr: _, dst_addr: _ } => 12,
      ProxyAddr::Ipv6Addr{ src_addr: _, dst_addr: _ } => 36,
      ProxyAddr::UnixAddr{ src_addr: _, dst_addr: _ } => 216,
      ProxyAddr::AfUnspec => 0,
    }
  }

  fn into_bytes(&self, buf: &mut Vec<u8>) {
    match *self {
      ProxyAddr::Ipv4Addr{ src_addr, dst_addr } => {
        buf.extend_from_slice(&src_addr.ip().octets());
        buf.extend_from_slice(&dst_addr.ip().octets());
        buf.extend_from_slice(&u16_to_array_of_u8(src_addr.port()));
        buf.extend_from_slice(&u16_to_array_of_u8(dst_addr.port()));
      },
      ProxyAddr::Ipv6Addr{ src_addr, dst_addr } => {
        buf.extend_from_slice(&src_addr.ip().octets());
        buf.extend_from_slice(&dst_addr.ip().octets());
        buf.extend_from_slice(&u16_to_array_of_u8(src_addr.port()));
        buf.extend_from_slice(&u16_to_array_of_u8(dst_addr.port()));
      },
      ProxyAddr::UnixAddr{ src_addr, dst_addr } => {
        buf.extend_from_slice(&src_addr);
        buf.extend_from_slice(&dst_addr);
      },
      ProxyAddr::AfUnspec => {},
    };
  }
}

fn get_family(addr: &ProxyAddr) -> u8 {
  match addr {
    &ProxyAddr::Ipv4Addr{ src_addr: _, dst_addr: _ } => 0x10 | 0x01, // AF_INET  = 1 + STREAM = 1
    &ProxyAddr::Ipv6Addr{ src_addr: _, dst_addr: _ } => 0x20 | 0x01, // AF_INET6 = 2 + STREAM = 1
    &ProxyAddr::UnixAddr{ src_addr: _, dst_addr: _ } => 0x30 | 0x01, // AF_UNIX  = 3 + STREAM = 1
    &ProxyAddr::AfUnspec => 0x00, // AF_UNSPEC + UNSPEC
  }
}

fn u16_to_array_of_u8(x: u16) -> [u8; 2] {
    let b1 : u8 = ((x >> 8) & 0xff) as u8;
    let b2 : u8 = (x & 0xff) as u8;
    return [b1, b2]
}

#[cfg(test)]
mod testV2 {

  use super::*;
  use std::net::{Ipv4Addr, Ipv6Addr};

  #[test]
  fn test_u16_to_array_of_u8() {
    let val_u16: u16 = 65534;
    let expected = [0xff, 0xfe];
    assert_eq!(expected, u16_to_array_of_u8(val_u16));
  }

  #[test]
  fn test_deserialize_tcp_ipv4_proxy_protocol_header() {
    let src_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(125, 25, 10, 1)), 8080);
    let dst_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 4, 5, 8)), 4200);

    let header = HeaderV2::new(src_addr, dst_addr);
    let expected = &[
      0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A, // MAGIC header
      0x20,                                                                   // Version 2 and command LOCAL
      0x11,                                                                   // family AF_UNIX with IPv4
      0x00, 0x0C,                                                             // address sizes = 12
      0x7D, 0x19, 0x0A, 0x01,                                                 // source address
      0x0A, 0x04, 0x05, 0x08,                                                 // destination address
      0x1F, 0x90,                                                             // source port
      0x10, 0x68,                                                             // destination port
    ];

    assert_eq!(expected, &header.into_bytes()[..]);
  }

  #[test]
  fn test_deserialize_tcp_ipv6_proxy_protocol_header() {
    let src_addr = SocketAddr::new(IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1)), 8080);
    let dst_addr = SocketAddr::new(IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1)), 4200);

    let header = HeaderV2::new(src_addr, dst_addr);
    let expected = vec![
      0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A,                         // MAGIC header
      0x20,                                                                                // Version 2 and command LOCAL
      0x21,                                                                                // family AF_UNIX with IPv6
      0x00, 0x24,                                                                                // address sizes = 36
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, // source address
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, // destination address
      0x1F, 0x90,                                                                                // source port
      0x10, 0x68,                                                                                // destination port
    ];

    assert_eq!(&expected[..], &header.into_bytes()[..]);
  }
}<|endoftext|>