<commit_before>//! Builder type for [`TwitterStream`](crate::TwitterStream).

pub use http::Method as RequestMethod;
pub use http::StatusCode;
pub use http::Uri;

use std::borrow::{Borrow, Cow};
use std::fmt::{self, Formatter};

use http::header::{HeaderValue, AUTHORIZATION, CONTENT_LENGTH, CONTENT_TYPE};
use http::Request;

use crate::service::HttpService;
use crate::token::Token;
use crate::util::fmt_join;
use crate::FutureTwitterStream;

/// A builder for [`TwitterStream`](crate::TwitterStream).
///
/// ## Example
///
/// ```rust,no_run
/// use futures::prelude::*;
/// use twitter_stream::{builder::BoundingBox, Token};
///
/// # #[tokio::main]
/// # async fn main() {
/// let token = Token::new("consumer_key", "consumer_secret", "access_key", "access_secret");
///
/// const TOKYO_WARDS: &'static [BoundingBox] = &[BoundingBox::new((139.56, 35.53), (139.92, 35.82))];
///
/// // Prints geolocated English Tweets associated with the special wards of Tokyo.
/// twitter_stream::Builder::filter(token)
///     .locations(TOKYO_WARDS)
///     .language("en")
///     .listen()
///     .try_flatten_stream()
///     .try_for_each(|json| {
///         println!("{}", json);
///         future::ok(())
///     })
///     .await
///     .unwrap();
/// # }
/// ```
#[derive(Clone, Debug)]
pub struct Builder<'a, T = Token> {
    method: RequestMethod,
    endpoint: Uri,
    token: T,
    parameters: Parameters<'a>,
}

/// Parameters to the Streaming API.
#[derive(Clone, Debug, Default, oauth::Authorize)]
struct Parameters<'a> {
    #[oauth1(skip_if = "not")]
    stall_warnings: bool,
    filter_level: Option<FilterLevel>,
    #[oauth1(skip_if = "str::is_empty")]
    language: Cow<'a, str>,
    #[oauth1(encoded, fmt = "fmt_follow", skip_if = "<[_]>::is_empty")]
    follow: Cow<'a, [u64]>,
    #[oauth1(skip_if = "str::is_empty")]
    track: Cow<'a, str>,
    #[oauth1(encoded, fmt = "fmt_locations", skip_if = "<[_]>::is_empty")]
    #[allow(clippy::type_complexity)]
    locations: Cow<'a, [BoundingBox]>,
    #[oauth1(encoded)]
    count: Option<i32>,
}

/// A `BoundingBox` is a rectangular area on the globe specified by coordinates of
/// the southwest and northeast edges in decimal degrees.
#[derive(Clone, Copy, Debug, PartialEq)]
#[repr(C)]
pub struct BoundingBox {
    /// Longitude of the west side of the bounding box.
    pub west_longitude: f64,
    /// Latitude of the south side of the bounding box.
    pub south_latitude: f64,
    /// Longitude of the east side of the bounding box.
    pub east_longitude: f64,
    /// Latitude of the north side of the bounding box.
    pub north_latitude: f64,
}

str_enum! {
    /// Represents the [`filter_level`] parameter in API requests.
    ///
    /// [`filter_level`]: https://developer.twitter.com/en/docs/tweets/filter-realtime/guides/basic-stream-parameters#filter-level
    #[derive(Clone, Debug, PartialEq, Hash, Eq)]
    pub enum FilterLevel {
        /// `"none"`
        None = "none",
        /// `"low"`
        Low = "low",
        /// `"medium"`
        Medium = "medium",
    }
}

impl<'a, C, A> Builder<'a, Token<C, A>>
where
    C: Borrow<str>,
    A: Borrow<str>,
{
    /// Create a builder for `POST statuses/filter` endpoint.
    ///
    /// See the [Twitter Developer Documentation][1] for more information.
    ///
    /// [1]: https://dev.twitter.com/streaming/reference/post/statuses/filter
    pub fn filter(token: Token<C, A>) -> Self {
        const URI: &str = "https://stream.twitter.com/1.1/statuses/filter.json";
        Self::custom(RequestMethod::POST, Uri::from_static(URI), token)
    }

    /// Create a builder for `GET statuses/sample` endpoint.
    ///
    /// See the [Twitter Developer Documentation][1] for more information.
    ///
    /// [1]: https://dev.twitter.com/streaming/reference/get/statuses/sample
    pub fn sample(token: Token<C, A>) -> Self {
        const URI: &str = "https://stream.twitter.com/1.1/statuses/sample.json";
        Self::custom(RequestMethod::GET, Uri::from_static(URI), token)
    }

    /// Constructs a builder for a Stream at a custom endpoint.
    pub fn custom(method: RequestMethod, endpoint: Uri, token: Token<C, A>) -> Self {
        Self {
            method,
            endpoint,
            token,
            parameters: Parameters::default(),
        }
    }

    /// Start listening on the Streaming API endpoint, returning a `Future` which resolves
    /// to a `Stream` yielding JSON messages from the API.
    ///
    /// # Panics
    ///
    /// This will panic if the underlying HTTPS connector failed to initialize.
    #[cfg(feature = "hyper")]
    #[cfg_attr(docsrs, doc(cfg(feature = "hyper")))]
    pub fn listen(&self) -> crate::hyper::FutureTwitterStream {
        let conn = hyper_tls::HttpsConnector::new();
        self.listen_with_client(hyper_pkg::Client::builder().build::<_, hyper_pkg::Body>(conn))
    }

    /// Same as [`listen`](Builder::listen) except that it uses `client` to make HTTP request
    /// to the endpoint.
    ///
    /// # Panics
    ///
    /// This will call `<S as Service>::call` without checking for `<S as Service>::poll_ready`
    /// and may cause a panic if `client` is not ready to send an HTTP request yet.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use tower::ServiceExt;
    ///
    /// # async fn doc() {
    /// # let mut client = hyper_pkg::Client::new();
    /// # let token = twitter_stream::Token::new("", "", "", "");
    /// client.ready().await; // Ensure that the client is ready to send a request.
    /// let stream = twitter_stream::Builder::sample(token)
    ///     .listen_with_client(&mut client)
    ///     .await
    ///     .unwrap();
    /// # }
    /// ```
    pub fn listen_with_client<S, B>(&self, mut client: S) -> FutureTwitterStream<S::Future>
    where
        S: HttpService<B>,
        B: Default + From<Vec<u8>>,
    {
        let req = Request::builder().method(self.method.clone());

        #[cfg(feature = "gzip")]
        let req = req.header(
            http::header::ACCEPT_ENCODING,
            HeaderValue::from_static("gzip"),
        );

        let mut oauth = oauth::Builder::new(self.token.client.as_ref(), oauth::HmacSha1);
        oauth.token(self.token.token.as_ref());
        let req = if RequestMethod::POST == self.method {
            let oauth::Request {
                authorization,
                data,
            } = oauth.post_form(&self.endpoint, &self.parameters);

            req.uri(self.endpoint.clone())
                .header(AUTHORIZATION, authorization)
                .header(
                    CONTENT_TYPE,
                    HeaderValue::from_static("application/x-www-form-urlencoded"),
                )
                .header(CONTENT_LENGTH, data.len())
                .body(data.into_bytes().into())
                .unwrap()
        } else {
            let oauth::Request {
                authorization,
                data: uri,
            } = oauth.build(self.method.as_ref(), &self.endpoint, &self.parameters);

            req.uri(uri)
                .header(AUTHORIZATION, authorization)
                .body(B::default())
                .unwrap()
        };

        let response = client.call(req);
        FutureTwitterStream { response }
    }
}

impl<'a, C, A> Builder<'a, Token<C, A>> {
    /// Reset the HTTP request method to be used when connecting to the server.
    pub fn method(&mut self, method: RequestMethod) -> &mut Self {
        self.method = method;
        self
    }

    /// Reset the API endpoint URI to be connected.
    pub fn endpoint(&mut self, endpoint: Uri) -> &mut Self {
        self.endpoint = endpoint;
        self
    }

    /// Reset the token to be used to log into Twitter.
    pub fn token(&mut self, token: Token<C, A>) -> &mut Self {
        self.token = token;
        self
    }

    /// Set whether to receive messages when in danger of being disconnected.
    ///
    /// See the [Twitter Developer Documentation][1] for more information.
    ///
    /// [1]: https://developer.twitter.com/en/docs/tweets/filter-realtime/guides/basic-stream-parameters#stall-warnings
    pub fn stall_warnings(&mut self, stall_warnings: bool) -> &mut Self {
        self.parameters.stall_warnings = stall_warnings;
        self
    }

    /// Set the minimum `filter_level` Tweet attribute to receive.
    /// The default is `FilterLevel::None`.
    ///
    /// See the [Twitter Developer Documentation][1] for more information.
    ///
    /// [1]: https://developer.twitter.com/en/docs/tweets/filter-realtime/guides/basic-stream-parameters#filter-level
    pub fn filter_level(&mut self, filter_level: impl Into<Option<FilterLevel>>) -> &mut Self {
        self.parameters.filter_level = filter_level.into();
        self
    }

    /// Set a comma-separated language identifiers to receive Tweets
    /// written in the specified languages only.
    ///
    /// Setting an empty string will unset this parameter.
    ///
    /// See the [Twitter Developer Documentation][1] for more information.
    ///
    /// [1]: https://developer.twitter.com/en/docs/tweets/filter-realtime/guides/basic-stream-parameters#language
    pub fn language(&mut self, language: impl Into<Cow<'a, str>>) -> &mut Self {
        self.parameters.language = language.into();
        self
    }

    /// Set a list of user IDs to receive Tweets from the specified users.
    ///
    /// Setting an empty slice will unset this parameter.
    ///
    /// This parameter is only available for the `filter` streams.
    ///
    /// See the [Twitter Developer Documentation][1] for more information.
    ///
    /// [1]: https://developer.twitter.com/en/docs/tweets/filter-realtime/guides/basic-stream-parameters#follow
    pub fn follow(&mut self, follow: impl Into<Cow<'a, [u64]>>) -> &mut Self {
        self.parameters.follow = follow.into();
        self
    }

    /// A comma separated list of phrases to filter Tweets by.
    ///
    /// Setting an empty string will unset this parameter.
    ///
    /// This parameter is only available for the `filter` streams.
    ///
    /// See the [Twitter Developer Documentation][1] for more information.
    ///
    /// [1]: https://developer.twitter.com/en/docs/tweets/filter-realtime/guides/basic-stream-parameters#track
    pub fn track(&mut self, track: impl Into<Cow<'a, str>>) -> &mut Self {
        self.parameters.track = track.into();
        self
    }

    /// Set a list of bounding boxes to filter Tweets by.
    ///
    /// Setting an empty slice will unset this parameter.
    ///
    /// This parameter is only available for the `filter` streams.
    ///
    /// See [`BoundingBox`](struct.BoundingBox.html) and
    /// the [Twitter Developer Documentation][1] for more information.
    ///
    /// [1]: https://developer.twitter.com/en/docs/tweets/filter-realtime/guides/basic-stream-parameters#locations
    pub fn locations(&mut self, locations: impl Into<Cow<'a, [BoundingBox]>>) -> &mut Self {
        self.parameters.locations = locations.into();
        self
    }

    /// The `count` parameter.
    /// This parameter requires elevated access to use.
    ///
    /// See the [Twitter Developer Documentation][1] for more information.
    ///
    /// [1]: https://developer.twitter.com/en/docs/tweets/filter-realtime/guides/basic-stream-parameters#count
    pub fn count(&mut self, count: impl Into<Option<i32>>) -> &mut Self {
        self.parameters.count = count.into();
        self
    }
}

impl BoundingBox {
    /// Creates a `BoundingBox` with two `(longitude, latitude)` pairs.
    ///
    /// The first argument specifies the southwest edge and the second specifies the northeast edge.
    ///
    /// # Example
    ///
    /// ```rust
    /// use twitter_stream::builder::BoundingBox;
    ///
    /// // Examples taken from Twitter's documentation.
    /// BoundingBox::new((-122.75, 36.8), (-121.75, 37.8)); // San Francisco
    /// BoundingBox::new((-74.0, 40.0), (-73.0, 41.0)); // New York City
    /// ```
    pub const fn new(
        (west_longitude, south_latitude): (f64, f64),
        (east_longitude, north_latitude): (f64, f64),
    ) -> Self {
        BoundingBox {
            west_longitude,
            south_latitude,
            east_longitude,
            north_latitude,
        }
    }
}

impl From<(f64, f64, f64, f64)> for BoundingBox {
    fn from(
        (west_longitude, south_latitude, east_longitude, north_latitude): (f64, f64, f64, f64),
    ) -> Self {
        BoundingBox {
            west_longitude,
            south_latitude,
            east_longitude,
            north_latitude,
        }
    }
}

impl From<((f64, f64), (f64, f64))> for BoundingBox {
    fn from(
        ((west_longitude, south_latitude), (east_longitude, north_latitude)): (
            (f64, f64),
            (f64, f64),
        ),
    ) -> Self {
        BoundingBox {
            west_longitude,
            south_latitude,
            east_longitude,
            north_latitude,
        }
    }
}

impl std::default::Default for FilterLevel {
    fn default() -> Self {
        FilterLevel::None
    }
}

impl std::fmt::Display for FilterLevel {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        AsRef::<str>::as_ref(self).fmt(f)
    }
}

const COMMA: &str = "%2C";

fn fmt_follow(ids: &[u64], f: &mut Formatter<'_>) -> fmt::Result {
    fmt_join(ids, COMMA, f)
}

fn fmt_locations(locs: &[BoundingBox], f: &mut Formatter<'_>) -> fmt::Result {
    use std::mem::size_of;
    use std::slice;

    use static_assertions::const_assert;

    let locs: &[f64] = unsafe {
        // Safety:
        // `BoundingBox` is defined with the `#[repr(C)]` attribute so it is sound to interpret
        // the struct as a `[f64; 4]` and also the fields are guaranteed to be placed in
        // `[west_longitude, south_latitude, east_longitude, north_latitude]` order.
        // We are checking the size of `BoundingBox` here just to be sure.
        const_assert!(size_of::<BoundingBox>() == 4 * size_of::<f64>());
        let ptr: *const BoundingBox = locs.as_ptr();
        let n = 4 * <[BoundingBox]>::len(locs);
        slice::from_raw_parts(ptr as *const f64, n)
    };

    fmt_join(locs, COMMA, f)
}

#[allow(clippy::trivially_copy_pass_by_ref)]
fn not(p: &bool) -> bool {
    !p
}
<commit_msg>Remove trait bound `B: Default` from `Builder::listen_with_client`<commit_after>//! Builder type for [`TwitterStream`](crate::TwitterStream).

pub use http::Method as RequestMethod;
pub use http::StatusCode;
pub use http::Uri;

use std::borrow::{Borrow, Cow};
use std::fmt::{self, Formatter};

use http::header::{HeaderValue, AUTHORIZATION, CONTENT_LENGTH, CONTENT_TYPE};
use http::Request;

use crate::service::HttpService;
use crate::token::Token;
use crate::util::fmt_join;
use crate::FutureTwitterStream;

/// A builder for [`TwitterStream`](crate::TwitterStream).
///
/// ## Example
///
/// ```rust,no_run
/// use futures::prelude::*;
/// use twitter_stream::{builder::BoundingBox, Token};
///
/// # #[tokio::main]
/// # async fn main() {
/// let token = Token::new("consumer_key", "consumer_secret", "access_key", "access_secret");
///
/// const TOKYO_WARDS: &'static [BoundingBox] = &[BoundingBox::new((139.56, 35.53), (139.92, 35.82))];
///
/// // Prints geolocated English Tweets associated with the special wards of Tokyo.
/// twitter_stream::Builder::filter(token)
///     .locations(TOKYO_WARDS)
///     .language("en")
///     .listen()
///     .try_flatten_stream()
///     .try_for_each(|json| {
///         println!("{}", json);
///         future::ok(())
///     })
///     .await
///     .unwrap();
/// # }
/// ```
#[derive(Clone, Debug)]
pub struct Builder<'a, T = Token> {
    method: RequestMethod,
    endpoint: Uri,
    token: T,
    parameters: Parameters<'a>,
}

/// Parameters to the Streaming API.
#[derive(Clone, Debug, Default, oauth::Authorize)]
struct Parameters<'a> {
    #[oauth1(skip_if = "not")]
    stall_warnings: bool,
    filter_level: Option<FilterLevel>,
    #[oauth1(skip_if = "str::is_empty")]
    language: Cow<'a, str>,
    #[oauth1(encoded, fmt = "fmt_follow", skip_if = "<[_]>::is_empty")]
    follow: Cow<'a, [u64]>,
    #[oauth1(skip_if = "str::is_empty")]
    track: Cow<'a, str>,
    #[oauth1(encoded, fmt = "fmt_locations", skip_if = "<[_]>::is_empty")]
    #[allow(clippy::type_complexity)]
    locations: Cow<'a, [BoundingBox]>,
    #[oauth1(encoded)]
    count: Option<i32>,
}

/// A `BoundingBox` is a rectangular area on the globe specified by coordinates of
/// the southwest and northeast edges in decimal degrees.
#[derive(Clone, Copy, Debug, PartialEq)]
#[repr(C)]
pub struct BoundingBox {
    /// Longitude of the west side of the bounding box.
    pub west_longitude: f64,
    /// Latitude of the south side of the bounding box.
    pub south_latitude: f64,
    /// Longitude of the east side of the bounding box.
    pub east_longitude: f64,
    /// Latitude of the north side of the bounding box.
    pub north_latitude: f64,
}

str_enum! {
    /// Represents the [`filter_level`] parameter in API requests.
    ///
    /// [`filter_level`]: https://developer.twitter.com/en/docs/tweets/filter-realtime/guides/basic-stream-parameters#filter-level
    #[derive(Clone, Debug, PartialEq, Hash, Eq)]
    pub enum FilterLevel {
        /// `"none"`
        None = "none",
        /// `"low"`
        Low = "low",
        /// `"medium"`
        Medium = "medium",
    }
}

impl<'a, C, A> Builder<'a, Token<C, A>>
where
    C: Borrow<str>,
    A: Borrow<str>,
{
    /// Create a builder for `POST statuses/filter` endpoint.
    ///
    /// See the [Twitter Developer Documentation][1] for more information.
    ///
    /// [1]: https://dev.twitter.com/streaming/reference/post/statuses/filter
    pub fn filter(token: Token<C, A>) -> Self {
        const URI: &str = "https://stream.twitter.com/1.1/statuses/filter.json";
        Self::custom(RequestMethod::POST, Uri::from_static(URI), token)
    }

    /// Create a builder for `GET statuses/sample` endpoint.
    ///
    /// See the [Twitter Developer Documentation][1] for more information.
    ///
    /// [1]: https://dev.twitter.com/streaming/reference/get/statuses/sample
    pub fn sample(token: Token<C, A>) -> Self {
        const URI: &str = "https://stream.twitter.com/1.1/statuses/sample.json";
        Self::custom(RequestMethod::GET, Uri::from_static(URI), token)
    }

    /// Constructs a builder for a Stream at a custom endpoint.
    pub fn custom(method: RequestMethod, endpoint: Uri, token: Token<C, A>) -> Self {
        Self {
            method,
            endpoint,
            token,
            parameters: Parameters::default(),
        }
    }

    /// Start listening on the Streaming API endpoint, returning a `Future` which resolves
    /// to a `Stream` yielding JSON messages from the API.
    ///
    /// # Panics
    ///
    /// This will panic if the underlying HTTPS connector failed to initialize.
    #[cfg(feature = "hyper")]
    #[cfg_attr(docsrs, doc(cfg(feature = "hyper")))]
    pub fn listen(&self) -> crate::hyper::FutureTwitterStream {
        let conn = hyper_tls::HttpsConnector::new();
        self.listen_with_client(hyper_pkg::Client::builder().build::<_, hyper_pkg::Body>(conn))
    }

    /// Same as [`listen`](Builder::listen) except that it uses `client` to make HTTP request
    /// to the endpoint.
    ///
    /// # Panics
    ///
    /// This will call `<S as Service>::call` without checking for `<S as Service>::poll_ready`
    /// and may cause a panic if `client` is not ready to send an HTTP request yet.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use tower::ServiceExt;
    ///
    /// # async fn doc() {
    /// # let mut client = hyper_pkg::Client::new();
    /// # let token = twitter_stream::Token::new("", "", "", "");
    /// client.ready().await; // Ensure that the client is ready to send a request.
    /// let stream = twitter_stream::Builder::sample(token)
    ///     .listen_with_client(&mut client)
    ///     .await
    ///     .unwrap();
    /// # }
    /// ```
    pub fn listen_with_client<S, B>(&self, mut client: S) -> FutureTwitterStream<S::Future>
    where
        S: HttpService<B>,
        B: From<Vec<u8>>,
    {
        let req = prepare_request(
            &self.method,
            &self.endpoint,
            self.token.as_ref(),
            &self.parameters,
        );
        let response = client.call(req.map(Into::into));
        FutureTwitterStream { response }
    }
}

impl<'a, C, A> Builder<'a, Token<C, A>> {
    /// Reset the HTTP request method to be used when connecting to the server.
    pub fn method(&mut self, method: RequestMethod) -> &mut Self {
        self.method = method;
        self
    }

    /// Reset the API endpoint URI to be connected.
    pub fn endpoint(&mut self, endpoint: Uri) -> &mut Self {
        self.endpoint = endpoint;
        self
    }

    /// Reset the token to be used to log into Twitter.
    pub fn token(&mut self, token: Token<C, A>) -> &mut Self {
        self.token = token;
        self
    }

    /// Set whether to receive messages when in danger of being disconnected.
    ///
    /// See the [Twitter Developer Documentation][1] for more information.
    ///
    /// [1]: https://developer.twitter.com/en/docs/tweets/filter-realtime/guides/basic-stream-parameters#stall-warnings
    pub fn stall_warnings(&mut self, stall_warnings: bool) -> &mut Self {
        self.parameters.stall_warnings = stall_warnings;
        self
    }

    /// Set the minimum `filter_level` Tweet attribute to receive.
    /// The default is `FilterLevel::None`.
    ///
    /// See the [Twitter Developer Documentation][1] for more information.
    ///
    /// [1]: https://developer.twitter.com/en/docs/tweets/filter-realtime/guides/basic-stream-parameters#filter-level
    pub fn filter_level(&mut self, filter_level: impl Into<Option<FilterLevel>>) -> &mut Self {
        self.parameters.filter_level = filter_level.into();
        self
    }

    /// Set a comma-separated language identifiers to receive Tweets
    /// written in the specified languages only.
    ///
    /// Setting an empty string will unset this parameter.
    ///
    /// See the [Twitter Developer Documentation][1] for more information.
    ///
    /// [1]: https://developer.twitter.com/en/docs/tweets/filter-realtime/guides/basic-stream-parameters#language
    pub fn language(&mut self, language: impl Into<Cow<'a, str>>) -> &mut Self {
        self.parameters.language = language.into();
        self
    }

    /// Set a list of user IDs to receive Tweets from the specified users.
    ///
    /// Setting an empty slice will unset this parameter.
    ///
    /// This parameter is only available for the `filter` streams.
    ///
    /// See the [Twitter Developer Documentation][1] for more information.
    ///
    /// [1]: https://developer.twitter.com/en/docs/tweets/filter-realtime/guides/basic-stream-parameters#follow
    pub fn follow(&mut self, follow: impl Into<Cow<'a, [u64]>>) -> &mut Self {
        self.parameters.follow = follow.into();
        self
    }

    /// A comma separated list of phrases to filter Tweets by.
    ///
    /// Setting an empty string will unset this parameter.
    ///
    /// This parameter is only available for the `filter` streams.
    ///
    /// See the [Twitter Developer Documentation][1] for more information.
    ///
    /// [1]: https://developer.twitter.com/en/docs/tweets/filter-realtime/guides/basic-stream-parameters#track
    pub fn track(&mut self, track: impl Into<Cow<'a, str>>) -> &mut Self {
        self.parameters.track = track.into();
        self
    }

    /// Set a list of bounding boxes to filter Tweets by.
    ///
    /// Setting an empty slice will unset this parameter.
    ///
    /// This parameter is only available for the `filter` streams.
    ///
    /// See [`BoundingBox`](struct.BoundingBox.html) and
    /// the [Twitter Developer Documentation][1] for more information.
    ///
    /// [1]: https://developer.twitter.com/en/docs/tweets/filter-realtime/guides/basic-stream-parameters#locations
    pub fn locations(&mut self, locations: impl Into<Cow<'a, [BoundingBox]>>) -> &mut Self {
        self.parameters.locations = locations.into();
        self
    }

    /// The `count` parameter.
    /// This parameter requires elevated access to use.
    ///
    /// See the [Twitter Developer Documentation][1] for more information.
    ///
    /// [1]: https://developer.twitter.com/en/docs/tweets/filter-realtime/guides/basic-stream-parameters#count
    pub fn count(&mut self, count: impl Into<Option<i32>>) -> &mut Self {
        self.parameters.count = count.into();
        self
    }
}

impl BoundingBox {
    /// Creates a `BoundingBox` with two `(longitude, latitude)` pairs.
    ///
    /// The first argument specifies the southwest edge and the second specifies the northeast edge.
    ///
    /// # Example
    ///
    /// ```rust
    /// use twitter_stream::builder::BoundingBox;
    ///
    /// // Examples taken from Twitter's documentation.
    /// BoundingBox::new((-122.75, 36.8), (-121.75, 37.8)); // San Francisco
    /// BoundingBox::new((-74.0, 40.0), (-73.0, 41.0)); // New York City
    /// ```
    pub const fn new(
        (west_longitude, south_latitude): (f64, f64),
        (east_longitude, north_latitude): (f64, f64),
    ) -> Self {
        BoundingBox {
            west_longitude,
            south_latitude,
            east_longitude,
            north_latitude,
        }
    }
}

impl From<(f64, f64, f64, f64)> for BoundingBox {
    fn from(
        (west_longitude, south_latitude, east_longitude, north_latitude): (f64, f64, f64, f64),
    ) -> Self {
        BoundingBox {
            west_longitude,
            south_latitude,
            east_longitude,
            north_latitude,
        }
    }
}

impl From<((f64, f64), (f64, f64))> for BoundingBox {
    fn from(
        ((west_longitude, south_latitude), (east_longitude, north_latitude)): (
            (f64, f64),
            (f64, f64),
        ),
    ) -> Self {
        BoundingBox {
            west_longitude,
            south_latitude,
            east_longitude,
            north_latitude,
        }
    }
}

impl std::default::Default for FilterLevel {
    fn default() -> Self {
        FilterLevel::None
    }
}

impl std::fmt::Display for FilterLevel {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        AsRef::<str>::as_ref(self).fmt(f)
    }
}

fn prepare_request(
    method: &RequestMethod,
    endpoint: &Uri,
    token: Token<&str, &str>,
    parameters: &Parameters,
) -> http::Request<Vec<u8>> {
    let req = Request::builder().method(method.clone());

    #[cfg(feature = "gzip")]
    let req = req.header(
        http::header::ACCEPT_ENCODING,
        HeaderValue::from_static("gzip"),
    );

    let mut oauth = oauth::Builder::new(token.client.as_ref(), oauth::HmacSha1);
    oauth.token(token.token.as_ref());

    if RequestMethod::POST == method {
        let oauth::Request {
            authorization,
            data,
        } = oauth.post_form(endpoint, parameters);

        req.uri(endpoint.clone())
            .header(AUTHORIZATION, authorization)
            .header(
                CONTENT_TYPE,
                HeaderValue::from_static("application/x-www-form-urlencoded"),
            )
            .header(CONTENT_LENGTH, data.len())
            .body(data.into_bytes())
            .unwrap()
    } else {
        let oauth::Request {
            authorization,
            data: uri,
        } = oauth.build(method.as_ref(), endpoint, parameters);

        req.uri(uri)
            .header(AUTHORIZATION, authorization)
            .body(Default::default())
            .unwrap()
    }
}

const COMMA: &str = "%2C";

fn fmt_follow(ids: &[u64], f: &mut Formatter<'_>) -> fmt::Result {
    fmt_join(ids, COMMA, f)
}

fn fmt_locations(locs: &[BoundingBox], f: &mut Formatter<'_>) -> fmt::Result {
    use std::mem::size_of;
    use std::slice;

    use static_assertions::const_assert;

    let locs: &[f64] = unsafe {
        // Safety:
        // `BoundingBox` is defined with the `#[repr(C)]` attribute so it is sound to interpret
        // the struct as a `[f64; 4]` and also the fields are guaranteed to be placed in
        // `[west_longitude, south_latitude, east_longitude, north_latitude]` order.
        // We are checking the size of `BoundingBox` here just to be sure.
        const_assert!(size_of::<BoundingBox>() == 4 * size_of::<f64>());
        let ptr: *const BoundingBox = locs.as_ptr();
        let n = 4 * <[BoundingBox]>::len(locs);
        slice::from_raw_parts(ptr as *const f64, n)
    };

    fmt_join(locs, COMMA, f)
}

#[allow(clippy::trivially_copy_pass_by_ref)]
fn not(p: &bool) -> bool {
    !p
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::net::SocketAddr;
use std::rc::{Rc, Weak};
use std::cell::RefCell;

use mioco::udp::UdpSocket;
use mioco::sync::mpsc::{channel, Receiver, Sender};
use bincode::rustc_serialize as bcode_rcs;
use sodiumoxide::crypto::secretbox as crypto_secretbox;

use ::Result;
use identity::{Identity, Extension, RemoteServer};
use packet::*;
use keys::*;
use boxes::*;
use nonces::*;

pub struct Stream {

}

impl Stream {
    pub fn new() -> Stream {
        Stream {}
    }
}

trait PacketProcessor {
    fn process_packet(&mut self, packet: Packet, sock: &mut UdpSocket, rem_addr: SocketAddr) -> Result<()>;
}

pub struct Demultiplexor {
    sock: UdpSocket,
    packet_processors: HashMap<Extension, Weak<RefCell<PacketProcessor>>>,
}

impl Demultiplexor {
    pub fn new(sock: UdpSocket) -> Demultiplexor {
        Demultiplexor {
            sock: sock,
            packet_processors: HashMap::new(),
        }
    }

    pub fn create_listener(&mut self, listener_id: Identity) -> Rc<RefCell<Listener>> {
        let listener_extension = listener_id.extension.clone();
        let listener = Rc::new(RefCell::new(Listener::new(listener_id)));

        self.packet_processors.insert(listener_extension, Rc::downgrade(&listener) as Weak<RefCell<PacketProcessor>>);

        listener
    }

    // pub fn create_client(&mut self, client_id: Identity, remote_id: RemoteServer) -> Rc<RefCell<ClientSock>> {
    //     let client_extension = client_id.extension.clone();
    //     let client = Rc::new(RefCell::new(ClientSock::new(client_id, remote_id)));
    //
    //     self.clients.insert(client_extension, Rc::downgrade(&client));
    //
    //     client
    // }

    pub fn listen(&mut self) -> Result<()> {
        let mut buf = [0u8; PACKET_MAX_SIZE];
        loop {
            let (recv_len, rem_addr) = try!(self.sock.recv(&mut buf));
            let packet: Packet = try!(bcode_rcs::decode(&buf[..recv_len]));

            if let Some(ref mut packet_processor) = {
                let dst_extension = packet.get_destination_extension();
                self.packet_processors.get_mut(dst_extension).and_then(|packet_processor_weak|
                    packet_processor_weak.upgrade()
                )
            } {
                try!(packet_processor.borrow_mut().process_packet(packet, &mut self.sock, rem_addr));
            }
        }
    }
}

pub struct ClientSock {
    // recv_rx: Receiver<Vec<u8>>,
    // recv_tx: Sender<Vec<u8>>,
}

impl ClientSock {
    pub fn new(my_id: Identity, remote_id: RemoteServer) -> ClientSock {
        // let (recv_rx, recv_tx) = channel();

        ClientSock {}
    }

    pub fn process(&mut self, packet: Packet, sock: &mut UdpSocket, rem_addr: SocketAddr) -> Result<()> {
        match packet {
            Packet::ServerMessage(server_msg_packet) => {
                try!(self.process_server_msg(server_msg_packet));
            },
            Packet::Cookie(cookie_packet) => {
            },
            _ => {
                debug!("Unvalid packet type");
            }
        }

        Ok(())
    }

    fn process_server_msg(&mut self, server_msg_packet: ServerMessagePacket) -> Result<()> {
        Err("Not implemented".into())
    }
}

pub struct ServerSocket {
    recv_rx: Receiver<Vec<u8>>,
    sock: UdpSocket,
    rem_addr: SocketAddr,
    client_extension: Extension,
    my_extension: Extension,
    precomputed_key: PrecomputedKey,
    next_send_nonce: Nonce8,
}

impl ServerSocket {
    pub fn recv(&mut self) -> Result<Vec<u8>> {
        self.recv_rx.recv().or(Err("Coudlnt read from recv channel".into()))
    }

    pub fn send(&mut self, msg: &[u8]) -> Result<usize> {
        let server_msg_packet: Packet = ServerMessagePacket {
            client_extension: self.client_extension.clone(),
            server_extension: self.my_extension.clone(),
            payload_box: ServerMessageBox::seal_precomputed(msg, &self.next_send_nonce, &self.precomputed_key),
        }.into();

        self.next_send_nonce.increment();

        let sent = try!(server_msg_packet.send(&mut self.sock, &self.rem_addr));

        if sent >= SERVER_MSG_PACKET_BASE_SIZE {
            Ok(sent - SERVER_MSG_PACKET_BASE_SIZE)
        } else {
            Err("Header trimmed".into())
        }
    }
}

pub struct Listener {
    my_extension: Extension,
    my_long_term_sk: server_long_term::SecretKey,
    minute_key: crypto_secretbox::Key,
    last_minute_key: crypto_secretbox::Key,
    conns: HashMap<(Extension, client_short_term::PublicKey), ServerConnection>,
    accept_tx: Sender<ServerSocket>,
    accept_rx: Receiver<ServerSocket>,
}

impl Listener {
    fn new(my_id: Identity) -> Listener {
        let (_, my_long_term_sk, my_extension) = my_id.as_server();
        let (accept_tx, accept_rx) = channel();

        Listener {
            my_extension: my_extension,
            my_long_term_sk: my_long_term_sk,
            minute_key: crypto_secretbox::gen_key(),
            last_minute_key: crypto_secretbox::gen_key(),
            conns: HashMap::new(),
            accept_tx: accept_tx,
            accept_rx: accept_rx,
        }
    }

    pub fn accept(&mut self) -> Result<ServerSocket> {
        self.accept_rx.recv().or(Err("Couldnt read accept channel".into()))
    }

    fn process_hello(&self, hello_packet: HelloPacket, sock: &mut UdpSocket, rem_addr: SocketAddr) -> Result<()> {
        let client_short_term_pk = hello_packet.client_short_term_pk;
        let conn_key = (hello_packet.client_extension.clone(), client_short_term_pk.clone());

        if self.conns.contains_key(&conn_key) {
            debug!("Hello packet received on open connection");
            return Ok(());
        }

        hello_packet.hello_box.open(&client_short_term_pk, &self.my_long_term_sk).unwrap();

        let (server_short_term_pk, server_short_term_sk) = server_short_term::gen_keypair();

        let cookie_packet: Packet = CookiePacket {
            client_extension: hello_packet.client_extension.clone(),
            server_extension: self.my_extension.clone(),
            cookie_box: PlainCookieBox {
                server_short_term_pk: server_short_term_pk,
                server_cookie: PlainCookie {
                    client_short_term_pk: client_short_term_pk.clone(),
                    server_short_term_sk: server_short_term_sk.clone(),
                }.seal(&CookieNonce::new_random(), &self.minute_key),
            }.seal(&Nonce16::new_random(), &client_short_term_pk, &self.my_long_term_sk),
        }.into();

        try!(cookie_packet.send(sock, &rem_addr));

        Ok(())
    }

    fn process_initiate(&mut self, initiate_packet: InitiatePacket, sock: &mut UdpSocket, rem_addr: SocketAddr) -> Result<()> {
        let conn_key = (initiate_packet.client_extension.clone(), initiate_packet.client_short_term_pk.clone());

        if self.conns.contains_key(&conn_key) {
            debug!("Recv'd INITIATE on already initiated connection");
            return Ok(());
        }

        let cookie = initiate_packet.server_cookie.open(&self.minute_key, &self.last_minute_key).unwrap();
        let precomputed_key = PrecomputedKey::precompute_at_server(&initiate_packet.client_short_term_pk, &cookie.server_short_term_sk);
        let (initiate_box, payload) = initiate_packet.initiate_box.open_precomputed_with_payload(&precomputed_key).unwrap();

        let client_long_term_pk = initiate_box.client_long_term_pk;
        let vouch = initiate_box.vouch.open(&client_long_term_pk, &self.my_long_term_sk).unwrap();

        if vouch.client_short_term_pk != initiate_packet.client_short_term_pk {
            return try!(Err("Invalid vouch"));
        }

        //TODO: Check if client_long_term_pk is accepted

        let (recv_tx, recv_rx) = channel();

        if let Some(msg) = payload {
            recv_tx.send(msg).unwrap();
        }

        let new_conn = ServerConnection {
            precomputed_key: precomputed_key.clone(),
            recv_tx: recv_tx,
            last_recv_nonce: initiate_packet.initiate_box.nonce.clone(),
        };

        self.conns.insert(conn_key, new_conn);
        info!("\"{}\" accepting connection \"{}\"@{}", &self.my_extension, &initiate_packet.client_extension, &rem_addr);

        let new_sock = ServerSocket {
            recv_rx: recv_rx,
            sock: sock.try_clone().unwrap(),
            rem_addr: rem_addr,
            client_extension: initiate_packet.client_extension.clone(),
            my_extension: self.my_extension.clone(),
            precomputed_key: precomputed_key,
            next_send_nonce: Nonce8::new_low(),
        };

        self.accept_tx.send(new_sock).or(Err("Couldnt read accept channel".into()))
    }

    fn process_client_msg(&mut self, client_msg_packet: ClientMessagePacket, _sock: &mut UdpSocket, _rem_addr: SocketAddr) -> Result<()> {
        let conn_key = (client_msg_packet.client_extension.clone(),
                        client_msg_packet.client_short_term_pk.clone());

        if let Some(server_conn) = self.conns.get_mut(&conn_key) {
            try!(server_conn.process_packet(client_msg_packet));
        }

        Ok(())
    }
}

impl PacketProcessor for Listener {
    fn process_packet(&mut self, packet: Packet, sock: &mut UdpSocket, rem_addr: SocketAddr) -> Result<()> {
        match packet {
            Packet::ClientMessage(client_msg_packet) => {
                try!(self.process_client_msg(client_msg_packet, sock, rem_addr));
            },
            Packet::Initiate(initiate_packet) => {
                try!(self.process_initiate(initiate_packet, sock, rem_addr));
            },
            Packet::Hello(hello_packet) => {
                try!(self.process_hello(hello_packet, sock, rem_addr));
            },
            _ => {
                debug!("Unvalid packet type");
            }
        }

        Ok(())
    }
}

struct ServerConnection {
    precomputed_key: PrecomputedKey,
    recv_tx: Sender<Vec<u8>>,
    last_recv_nonce: Nonce8,
}

impl ServerConnection {
    pub fn process_packet(&mut self, client_msg_packet: ClientMessagePacket) -> Result<()> {
        if client_msg_packet.payload_box.nonce <= self.last_recv_nonce {
            return Err("Bad nonce".into()); // TODO: Rework to allow some packet reordering
        }

        let msg = try!(client_msg_packet.payload_box.open_precomputed(&self.precomputed_key)
            .or(Err("Bad encrypted message")));

        self.last_recv_nonce = client_msg_packet.payload_box.nonce;

        try!(self.recv_tx.send(msg).or(Err("Couldnt write to recv channel")));

        Ok(())
    }
}
<commit_msg>impl PacketProcessor for ClientSock<commit_after>use std::collections::HashMap;
use std::net::SocketAddr;
use std::rc::{Rc, Weak};
use std::cell::RefCell;

use mioco::udp::UdpSocket;
use mioco::sync::mpsc::{channel, Receiver, Sender};
use bincode::rustc_serialize as bcode_rcs;
use sodiumoxide::crypto::secretbox as crypto_secretbox;

use ::Result;
use identity::{Identity, Extension, RemoteServer};
use packet::*;
use keys::*;
use boxes::*;
use nonces::*;

pub struct Stream {

}

impl Stream {
    pub fn new() -> Stream {
        Stream {}
    }
}

trait PacketProcessor {
    fn process_packet(&mut self, packet: Packet, sock: &mut UdpSocket, rem_addr: SocketAddr) -> Result<()>;
}

pub struct Demultiplexor {
    sock: UdpSocket,
    packet_processors: HashMap<Extension, Weak<RefCell<PacketProcessor>>>,
}

impl Demultiplexor {
    pub fn new(sock: UdpSocket) -> Demultiplexor {
        Demultiplexor {
            sock: sock,
            packet_processors: HashMap::new(),
        }
    }

    pub fn create_listener(&mut self, listener_id: Identity) -> Rc<RefCell<Listener>> {
        let listener_extension = listener_id.extension.clone();
        let listener = Rc::new(RefCell::new(Listener::new(listener_id)));

        self.packet_processors.insert(listener_extension, Rc::downgrade(&listener) as Weak<RefCell<PacketProcessor>>);

        listener
    }

    pub fn create_client(&mut self, client_id: Identity, remote_id: RemoteServer) -> Rc<RefCell<ClientSock>> {
        let client_extension = client_id.extension.clone();
        let client = Rc::new(RefCell::new(ClientSock::new(client_id, remote_id)));

        self.packet_processors.insert(client_extension, Rc::downgrade(&client) as Weak<RefCell<PacketProcessor>>);

        client
    }

    pub fn listen(&mut self) -> Result<()> {
        let mut buf = [0u8; PACKET_MAX_SIZE];
        loop {
            let (recv_len, rem_addr) = try!(self.sock.recv(&mut buf));
            let packet: Packet = try!(bcode_rcs::decode(&buf[..recv_len]));

            if let Some(ref mut packet_processor) = {
                let dst_extension = packet.get_destination_extension();
                self.packet_processors.get_mut(dst_extension).and_then(|packet_processor_weak|
                    packet_processor_weak.upgrade()
                )
            } {
                try!(packet_processor.borrow_mut().process_packet(packet, &mut self.sock, rem_addr));
            }
        }
    }
}

pub struct ClientSock {
    // recv_rx: Receiver<Vec<u8>>,
    // recv_tx: Sender<Vec<u8>>,
}

impl ClientSock {
    pub fn new(my_id: Identity, remote_id: RemoteServer) -> ClientSock {
        // let (recv_rx, recv_tx) = channel();

        ClientSock {}
    }

    fn process_server_msg(&mut self, server_msg_packet: ServerMessagePacket) -> Result<()> {
        Err("Not implemented".into())
    }
}

impl PacketProcessor for ClientSock {
    fn process_packet(&mut self, packet: Packet, sock: &mut UdpSocket, rem_addr: SocketAddr) -> Result<()> {
        match packet {
            Packet::ServerMessage(server_msg_packet) => {
                try!(self.process_server_msg(server_msg_packet));
            },
            Packet::Cookie(cookie_packet) => {
            },
            _ => {
                debug!("Unvalid packet type");
            }
        }

        Ok(())
    }
}

pub struct ServerSocket {
    recv_rx: Receiver<Vec<u8>>,
    sock: UdpSocket,
    rem_addr: SocketAddr,
    client_extension: Extension,
    my_extension: Extension,
    precomputed_key: PrecomputedKey,
    next_send_nonce: Nonce8,
}

impl ServerSocket {
    pub fn recv(&mut self) -> Result<Vec<u8>> {
        self.recv_rx.recv().or(Err("Coudlnt read from recv channel".into()))
    }

    pub fn send(&mut self, msg: &[u8]) -> Result<usize> {
        let server_msg_packet: Packet = ServerMessagePacket {
            client_extension: self.client_extension.clone(),
            server_extension: self.my_extension.clone(),
            payload_box: ServerMessageBox::seal_precomputed(msg, &self.next_send_nonce, &self.precomputed_key),
        }.into();

        self.next_send_nonce.increment();

        let sent = try!(server_msg_packet.send(&mut self.sock, &self.rem_addr));

        if sent >= SERVER_MSG_PACKET_BASE_SIZE {
            Ok(sent - SERVER_MSG_PACKET_BASE_SIZE)
        } else {
            Err("Header trimmed".into())
        }
    }
}

pub struct Listener {
    my_extension: Extension,
    my_long_term_sk: server_long_term::SecretKey,
    minute_key: crypto_secretbox::Key,
    last_minute_key: crypto_secretbox::Key,
    conns: HashMap<(Extension, client_short_term::PublicKey), ServerConnection>,
    accept_tx: Sender<ServerSocket>,
    accept_rx: Receiver<ServerSocket>,
}

impl Listener {
    fn new(my_id: Identity) -> Listener {
        let (_, my_long_term_sk, my_extension) = my_id.as_server();
        let (accept_tx, accept_rx) = channel();

        Listener {
            my_extension: my_extension,
            my_long_term_sk: my_long_term_sk,
            minute_key: crypto_secretbox::gen_key(),
            last_minute_key: crypto_secretbox::gen_key(),
            conns: HashMap::new(),
            accept_tx: accept_tx,
            accept_rx: accept_rx,
        }
    }

    pub fn accept(&mut self) -> Result<ServerSocket> {
        self.accept_rx.recv().or(Err("Couldnt read accept channel".into()))
    }

    fn process_hello(&self, hello_packet: HelloPacket, sock: &mut UdpSocket, rem_addr: SocketAddr) -> Result<()> {
        let client_short_term_pk = hello_packet.client_short_term_pk;
        let conn_key = (hello_packet.client_extension.clone(), client_short_term_pk.clone());

        if self.conns.contains_key(&conn_key) {
            debug!("Hello packet received on open connection");
            return Ok(());
        }

        hello_packet.hello_box.open(&client_short_term_pk, &self.my_long_term_sk).unwrap();

        let (server_short_term_pk, server_short_term_sk) = server_short_term::gen_keypair();

        let cookie_packet: Packet = CookiePacket {
            client_extension: hello_packet.client_extension.clone(),
            server_extension: self.my_extension.clone(),
            cookie_box: PlainCookieBox {
                server_short_term_pk: server_short_term_pk,
                server_cookie: PlainCookie {
                    client_short_term_pk: client_short_term_pk.clone(),
                    server_short_term_sk: server_short_term_sk.clone(),
                }.seal(&CookieNonce::new_random(), &self.minute_key),
            }.seal(&Nonce16::new_random(), &client_short_term_pk, &self.my_long_term_sk),
        }.into();

        try!(cookie_packet.send(sock, &rem_addr));

        Ok(())
    }

    fn process_initiate(&mut self, initiate_packet: InitiatePacket, sock: &mut UdpSocket, rem_addr: SocketAddr) -> Result<()> {
        let conn_key = (initiate_packet.client_extension.clone(), initiate_packet.client_short_term_pk.clone());

        if self.conns.contains_key(&conn_key) {
            debug!("Recv'd INITIATE on already initiated connection");
            return Ok(());
        }

        let cookie = initiate_packet.server_cookie.open(&self.minute_key, &self.last_minute_key).unwrap();
        let precomputed_key = PrecomputedKey::precompute_at_server(&initiate_packet.client_short_term_pk, &cookie.server_short_term_sk);
        let (initiate_box, payload) = initiate_packet.initiate_box.open_precomputed_with_payload(&precomputed_key).unwrap();

        let client_long_term_pk = initiate_box.client_long_term_pk;
        let vouch = initiate_box.vouch.open(&client_long_term_pk, &self.my_long_term_sk).unwrap();

        if vouch.client_short_term_pk != initiate_packet.client_short_term_pk {
            return try!(Err("Invalid vouch"));
        }

        //TODO: Check if client_long_term_pk is accepted

        let (recv_tx, recv_rx) = channel();

        if let Some(msg) = payload {
            recv_tx.send(msg).unwrap();
        }

        let new_conn = ServerConnection {
            precomputed_key: precomputed_key.clone(),
            recv_tx: recv_tx,
            last_recv_nonce: initiate_packet.initiate_box.nonce.clone(),
        };

        self.conns.insert(conn_key, new_conn);
        info!("\"{}\" accepting connection \"{}\"@{}", &self.my_extension, &initiate_packet.client_extension, &rem_addr);

        let new_sock = ServerSocket {
            recv_rx: recv_rx,
            sock: sock.try_clone().unwrap(),
            rem_addr: rem_addr,
            client_extension: initiate_packet.client_extension.clone(),
            my_extension: self.my_extension.clone(),
            precomputed_key: precomputed_key,
            next_send_nonce: Nonce8::new_low(),
        };

        self.accept_tx.send(new_sock).or(Err("Couldnt read accept channel".into()))
    }

    fn process_client_msg(&mut self, client_msg_packet: ClientMessagePacket, _sock: &mut UdpSocket, _rem_addr: SocketAddr) -> Result<()> {
        let conn_key = (client_msg_packet.client_extension.clone(),
                        client_msg_packet.client_short_term_pk.clone());

        if let Some(server_conn) = self.conns.get_mut(&conn_key) {
            try!(server_conn.process_packet(client_msg_packet));
        }

        Ok(())
    }
}

impl PacketProcessor for Listener {
    fn process_packet(&mut self, packet: Packet, sock: &mut UdpSocket, rem_addr: SocketAddr) -> Result<()> {
        match packet {
            Packet::ClientMessage(client_msg_packet) => {
                try!(self.process_client_msg(client_msg_packet, sock, rem_addr));
            },
            Packet::Initiate(initiate_packet) => {
                try!(self.process_initiate(initiate_packet, sock, rem_addr));
            },
            Packet::Hello(hello_packet) => {
                try!(self.process_hello(hello_packet, sock, rem_addr));
            },
            _ => {
                debug!("Unvalid packet type");
            }
        }

        Ok(())
    }
}

struct ServerConnection {
    precomputed_key: PrecomputedKey,
    recv_tx: Sender<Vec<u8>>,
    last_recv_nonce: Nonce8,
}

impl ServerConnection {
    pub fn process_packet(&mut self, client_msg_packet: ClientMessagePacket) -> Result<()> {
        if client_msg_packet.payload_box.nonce <= self.last_recv_nonce {
            return Err("Bad nonce".into()); // TODO: Rework to allow some packet reordering
        }

        let msg = try!(client_msg_packet.payload_box.open_precomputed(&self.precomputed_key)
            .or(Err("Bad encrypted message")));

        self.last_recv_nonce = client_msg_packet.payload_box.nonce;

        try!(self.recv_tx.send(msg).or(Err("Couldnt write to recv channel")));

        Ok(())
    }
}
<|endoftext|>


<commit_before>// mod channel

use rustc_serialize::json;
use client::{Client, ClientList};
use std::fmt;
use std::cmp;
use std::ops::Deref;
use std::str::FromStr;
use map::*;
use escaping::*;
use error;

/// Channel contains a TeamSpeak 3 channel.
/// # Example
/// ```
/// use sqlib::channel;
///
/// let channel = sqlib::channel::Channel::new(0, "test".to_string());
///
/// assert!(channel.is_empty());
/// assert_eq!("test".to_string(), format!("{}", channel));
/// ```
#[derive(Debug, Clone, RustcDecodable, RustcEncodable)]
pub struct Channel {
    pub cid: i64,
    pub channel_name: String,
    pub clients: Vec<Client>,
}

impl Default for Channel {
    fn default() -> Channel {
        Channel {
            cid: 0,
            channel_name: String::new(),
            clients: Vec::new(),
        }
    }
}

impl Channel {
    /// Create an empty Channel from a channel id and a name.
    pub fn new(channel_id: i64, name: String) -> Channel {
        let mut channel = Channel::default();
        channel.cid = channel_id;
        channel.channel_name = name;
        channel.unescape();
        channel
    }

    fn unescape(&mut self) {
        self.channel_name = unescape(&self.channel_name);
    }

    /// Create a new Channel from a given map.
    pub fn from_map(map: &StringMap) -> Channel {
        let mut channel = Channel::default();
        channel.mut_from_map(map);
        channel
    }

    /// Create a new Channel from a given Channel and a map.
    pub fn update_from_map(channel: Channel, map: &StringMap) -> Channel {
        let mut channel = channel.clone();
        channel.mut_from_map(map);
        channel
    }

    /// Mutate self from a given map.
    pub fn mut_from_map(&mut self, map: &StringMap) {
        update_from_map(map, "cid", &mut self.cid);
        update_from_map(map, "channel_name", &mut self.channel_name);
        self.unescape();
    }

    pub fn is_empty(&self) -> bool {
        self.clients.is_empty()
    }

    pub fn clients_len(&self) -> usize {
        self.clients.len()
    }

    pub fn add_client(&mut self, client: Client) {
        self.clients.push(client);
    }

    pub fn sort_clients(&mut self) {
        self.clients.sort()
    }

    /// Remove all Server Query Clients from the Channel.
    pub fn remove_sq_clients(&mut self) {
        let new_clients = self.clients
            .iter()
            .filter(|c| c.is_client())
            .map(|c| c.clone())
            .collect();
        self.clients = new_clients;
    }

    /// Creates a JSON String from self.
    pub fn as_json(&self) -> String {
        json::encode(self).unwrap_or(String::new())
    }
}

impl FromStr for Channel {
    type Err = error::Error;
    fn from_str(s: &str) -> error::Result<Self> {
        let map = to_map(s);
        Ok(Channel::from_map(&map))
    }
}

impl fmt::Display for Channel {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        try!(write!(f, "{}", self.channel_name));
        if !self.is_empty() {
            for client in self.clients.iter() {
                try!(write!(f, "\n  {}", client));
            }
        };
        Ok(())
    }
}

impl PartialEq for Channel {
    fn eq(&self, other: &Self) -> bool {
        self.cid.eq(&other.cid)
    }
}

impl Eq for Channel {}

impl PartialOrd for Channel {
    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
        self.cid.partial_cmp(&other.cid)
    }
}

impl Ord for Channel {
    fn cmp(&self, other: &Self) -> cmp::Ordering {
        self.cid.cmp(&other.cid)
    }
}

#[derive(Debug, Eq, PartialEq, Ord, PartialOrd, Clone, RustcDecodable, RustcEncodable)]
pub struct ChannelList(Vec<Channel>);

impl Default for ChannelList {
    fn default() -> ChannelList {
        ChannelList(Vec::new())
    }
}

impl ChannelList {
    pub fn into_inner(self) -> Vec<Channel> {
        let ChannelList(channel) = self;
        channel
    }

    pub fn from_maps(maps: &Vec<StringMap>) -> ChannelList {
        let mut vec = Vec::new();
        for map in maps.iter() {
            let channel = Channel::from_map(map);
            vec.push(channel);
        }
        ChannelList(vec)
    }

    pub fn insert_client(&mut self, client: &Client) {
        let cid = client.cid;
        for channel in self.as_mut().iter_mut() {
            if channel.cid == cid {
                channel.add_client(client.clone());
            }
        }
    }

    pub fn merge_clients(&mut self, clients: &ClientList) {
        for client in clients.iter() {
            self.insert_client(client);
        }
    }

    pub fn as_json(&self) -> String {
        json::encode(self.as_ref()).unwrap_or(String::new())
    }
}

impl Deref for ChannelList {
    type Target = Vec<Channel>;
    fn deref(&self) -> &Vec<Channel> {
        self.as_ref()
    }
}

impl AsRef<Vec<Channel>> for ChannelList {
    fn as_ref(&self) -> &Vec<Channel> {
        let &ChannelList(ref channel) = self;
        channel
    }
}

impl AsMut<Vec<Channel>> for ChannelList {
    fn as_mut(&mut self) -> &mut Vec<Channel> {
        let &mut ChannelList(ref mut channel) = self;
        channel
    }
}

impl FromStr for ChannelList {
    type Err = error::Error;
    fn from_str(s: &str) -> error::Result<Self> {
        let maps = s.split('|').map(to_map).collect();
        Ok(ChannelList::from_maps(&maps))
    }
}
<commit_msg>adding documentation and two trait implementation to sqlib::channel<commit_after>//! The channel module contains the Channel and ChannelList structs.
//!
//! The are representations of a TS3 channel and channellist.

use rustc_serialize::json;
use client::{Client, ClientList};
use std::fmt;
use std::cmp;
use std::ops::Deref;
use std::str::FromStr;
use map::*;
use escaping::*;
use error;

/// Channel contains a TeamSpeak 3 channel.
/// # Example
/// ```
/// use sqlib::channel;
///
/// let channel = sqlib::channel::Channel::new(0, "test".to_string());
///
/// assert!(channel.is_empty());
/// assert_eq!("test".to_string(), format!("{}", channel));
/// ```
#[derive(Debug, Clone, RustcDecodable, RustcEncodable)]
pub struct Channel {
    /// channel id
    pub cid: i64,
    /// channel name
    pub channel_name: String,
    /// A vector of clients, who are in the channel.
    pub clients: Vec<Client>,
}

impl Default for Channel {
    fn default() -> Channel {
        Channel {
            cid: 0,
            channel_name: String::new(),
            clients: Vec::new(),
        }
    }
}

impl Channel {
    /// Create an empty Channel from a channel id and a name.
    pub fn new(channel_id: i64, name: String) -> Channel {
        let mut channel = Channel::default();
        channel.cid = channel_id;
        channel.channel_name = name;
        channel.unescape();
        channel
    }

    fn unescape(&mut self) {
        self.channel_name = unescape(&self.channel_name);
    }

    /// Create a new Channel from a given map.
    pub fn from_map(map: &StringMap) -> Channel {
        let mut channel = Channel::default();
        channel.mut_from_map(map);
        channel
    }

    /// Create a new Channel from a given Channel and a map.
    pub fn update_from_map(channel: Channel, map: &StringMap) -> Channel {
        let mut channel = channel.clone();
        channel.mut_from_map(map);
        channel
    }

    /// Mutate self from a given map.
    pub fn mut_from_map(&mut self, map: &StringMap) {
        update_from_map(map, "cid", &mut self.cid);
        update_from_map(map, "channel_name", &mut self.channel_name);
        self.unescape();
    }

    pub fn is_empty(&self) -> bool {
        self.clients.is_empty()
    }

    pub fn clients_len(&self) -> usize {
        self.clients.len()
    }

    pub fn add_client(&mut self, client: Client) {
        self.clients.push(client);
    }

    pub fn sort_clients(&mut self) {
        self.clients.sort()
    }

    /// Remove all Server Query Clients from the Channel.
    pub fn remove_sq_clients(&mut self) {
        let new_clients = self.clients
            .iter()
            .filter(|c| c.is_client())
            .map(|c| c.clone())
            .collect();
        self.clients = new_clients;
    }

    /// Creates a JSON String from self.
    pub fn as_json(&self) -> String {
        json::encode(self).unwrap_or(String::new())
    }
}

impl FromStr for Channel {
    type Err = error::Error;
    fn from_str(s: &str) -> error::Result<Self> {
        let map = to_map(s);
        Ok(Channel::from_map(&map))
    }
}

impl fmt::Display for Channel {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        try!(write!(f, "{}", self.channel_name));
        if !self.is_empty() {
            for client in self.clients.iter() {
                try!(write!(f, "\n  {}", client));
            }
        };
        Ok(())
    }
}

impl PartialEq for Channel {
    fn eq(&self, other: &Self) -> bool {
        self.cid.eq(&other.cid)
    }
}

impl Eq for Channel {}

impl PartialOrd for Channel {
    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
        self.cid.partial_cmp(&other.cid)
    }
}

impl Ord for Channel {
    fn cmp(&self, other: &Self) -> cmp::Ordering {
        self.cid.cmp(&other.cid)
    }
}

/// ChannelList contains a list of Channels.
///
/// # Example
/// ```
/// use sqlib::channel::{Channel, ChannelList};
///
/// let channel1 = Channel::new(1, "test1".to_string());
/// let channel2 = Channel::new(2, "test2".to_string());
///
/// let channels = vec![channel1, channel2];
/// let channellist = ChannelList::from(channels.clone());
///
/// assert_eq!(&channels, channellist.as_ref());
///
/// let channels2: Vec<_> = channellist.into();
/// assert_eq!(channels, channels2);
/// ```
#[derive(Debug, Eq, PartialEq, Ord, PartialOrd, Clone, RustcDecodable, RustcEncodable)]
pub struct ChannelList(Vec<Channel>);

impl Default for ChannelList {
    fn default() -> ChannelList {
        ChannelList(Vec::new())
    }
}

impl ChannelList {
    /// Converts a ChannelList into a vector of Channel's.
    pub fn into_inner(self) -> Vec<Channel> {
        let ChannelList(channel) = self;
        channel
    }

    /// creates a ChannelList from a Vec of StringMaps
    pub fn from_maps(maps: &Vec<StringMap>) -> ChannelList {
        let mut vec = Vec::new();
        for map in maps.iter() {
            let channel = Channel::from_map(map);
            vec.push(channel);
        }
        ChannelList(vec)
    }

    pub fn insert_client(&mut self, client: &Client) {
        let cid = client.cid;
        for channel in self.as_mut().iter_mut() {
            if channel.cid == cid {
                channel.add_client(client.clone());
            }
        }
    }

    pub fn merge_clients(&mut self, clients: &ClientList) {
        for client in clients.iter() {
            self.insert_client(client);
        }
    }

    /// creates a JSON String from a ChannelList
    pub fn as_json(&self) -> String {
        json::encode(self.as_ref()).unwrap_or(String::new())
    }
}

impl Deref for ChannelList {
    type Target = Vec<Channel>;
    fn deref(&self) -> &Vec<Channel> {
        self.as_ref()
    }
}

impl AsRef<Vec<Channel>> for ChannelList {
    fn as_ref(&self) -> &Vec<Channel> {
        let &ChannelList(ref channel) = self;
        channel
    }
}

impl AsMut<Vec<Channel>> for ChannelList {
    fn as_mut(&mut self) -> &mut Vec<Channel> {
        let &mut ChannelList(ref mut channel) = self;
        channel
    }
}

impl FromStr for ChannelList {
    type Err = error::Error;
    fn from_str(s: &str) -> error::Result<Self> {
        let maps = s.split('|').map(to_map).collect();
        Ok(ChannelList::from_maps(&maps))
    }
}

impl From<Vec<Channel>> for ChannelList {
    fn from(channels: Vec<Channel>) -> Self {
        ChannelList(channels)
    }
}

impl From<ChannelList> for Vec<Channel> {
    fn from(channels: ChannelList) -> Self {
        channels.into_inner()
    }
}
<|endoftext|>


<commit_before>use async_trait::async_trait;
use cucumber::{then, when};
use geo::algorithm::haversine_distance::HaversineDistance;
use itertools::{EitherOrBoth::*, Itertools};
use mimir::adapters::secondary::elasticsearch::remote::connection_test_pool;
use mimir::domain::model::configuration;
use mimir::domain::ports::secondary::remote::Remote;
use std::cmp::Ordering;

use crate::error::Error;
use crate::state::{GlobalState, State, Step, StepStatus};
use mimir::adapters::primary::bragi::api::DEFAULT_LIMIT_RESULT_ES;
use mimir::adapters::primary::{
    common::coord::Coord, common::dsl::build_query, common::filters::Filters,
    common::settings::QuerySettings,
};
use mimir::adapters::secondary::elasticsearch::ElasticsearchStorageConfig;
use mimir::domain::{model::query::Query, ports::primary::search_documents::SearchDocuments};
use places::{addr::Addr, admin::Admin, poi::Poi};

// Search place

#[when(regex = r#"the user searches (.+) datatype for "(.*)" at (.+), (.+)$"#)]
async fn search(state: &mut GlobalState, places: String, query: String, lat: f32, lon: f32) {
    perform_search(state, places, query, Coord::new(lat, lon).into(), None).await;
}

#[when(regex = r#"the user searches (.+) datatype for "(.+)" with "(.+)" filters$"#)]
async fn search_with_zone_filters(
    state: &mut GlobalState,
    places: String,
    query: String,
    zone_types: String,
) {
    perform_search(state, places, query, None, Some(zone_types)).await;
}

#[when(regex = r#"the user searches (.+) datatype for "(.*)"$"#)]
async fn search_no_coord(state: &mut GlobalState, places: String, query: String) {
    perform_search(state, places, query, None, None).await;
}

async fn perform_search(
    state: &mut GlobalState,
    places: String,
    query: String,
    coord: Option<Coord>,
    zone_types: Option<String>,
) {
    let places = {
        if places == "all" {
            vec![configuration::root()]
        } else {
            places
                .split(',')
                .map(str::trim)
                .map(str::to_string)
                .map(|s| configuration::root_doctype(&s))
                .collect()
        }
    };

    let zone_types = match zone_types {
        Some(f) => Some(
            f.split(',')
                .map(str::trim)
                .map(str::to_string)
                .collect::<Vec<String>>(),
        ),
        _ => None,
    };

    let filters = Filters {
        coord,
        zone_types,
        ..Default::default()
    };

    state
        .execute(Search {
            places,
            query,
            filters,
            results: Vec::new(),
        })
        .await
        .expect("failed to search");
}

/// Perform a search in current Elasticsearch DB.
#[derive(Debug)]
pub struct Search {
    places: Vec<String>, // What kind of places we're searching
    query: String,       // The search string
    filters: Filters,    // Search filters
    results: Vec<serde_json::Value>,
}

#[async_trait(?Send)]
impl Step for Search {
    async fn execute(&mut self, _state: &State) -> Result<StepStatus, Error> {
        let client = connection_test_pool()
            .conn(ElasticsearchStorageConfig::default_testing())
            .await
            .expect("Could not establish connection to Elasticsearch");

        // Build ES query
        let dsl = build_query(
            &self.query,
            self.filters.clone(),
            &["fr"],
            &QuerySettings::default(),
        );

        // Fetch documents
        self.results = {
            client
                .search_documents(
                    self.places.clone(),
                    Query::QueryDSL(dsl),
                    DEFAULT_LIMIT_RESULT_ES,
                    None,
                )
                .await
                .unwrap()
        };

        Ok(StepStatus::Done)
    }
}

// Find document

#[then(regex = r#"he finds "(.+)" as the first result$"#)]
async fn find_id(state: &mut GlobalState, id: String) {
    state
        .execute(HasDocument { id, max_rank: 1 })
        .await
        .expect("failed to find document");
}

/// Check if given document is in the output.
///
/// It assumes that a Search has already been performed before.
pub struct HasDocument {
    id: String,
    max_rank: usize,
}

#[async_trait(?Send)]
impl Step for HasDocument {
    async fn execute(&mut self, state: &State) -> Result<StepStatus, Error> {
        let (search, _) = state
            .steps_for::<Search>()
            .next_back()
            .expect("the user must perform a search before checking results");

        let (rank, _) = (search.results.iter().enumerate())
            .find(|(_, doc)| {
                println!("{} vs {}", self.id, doc["id"].as_str().unwrap());
                self.id == doc["id"].as_str().unwrap()
            })
            .expect("document was not found in search results");

        assert!(rank < self.max_rank);
        Ok(StepStatus::Done)
    }
}

// Find Admin

#[then(regex = r#"he finds admin "(.*)", a "(.*)", in the first (.*) results$"#)]
async fn find_admin(state: &mut GlobalState, name: String, zone_type: String, limit: usize) {
    state
        .execute(HasAdmin {
            name: Some(name),
            zone_type: Some(zone_type),
            limit,
        })
        .await
        .expect("failed to find document");
}

/// Check if given admin is in the output.
///
/// It assumes that a Search has already been performed before.
pub struct HasAdmin {
    name: Option<String>,
    zone_type: Option<String>,
    limit: usize,
}

#[async_trait(?Send)]
impl Step for HasAdmin {
    async fn execute(&mut self, state: &State) -> Result<StepStatus, Error> {
        let (search, _) = state
            .steps_for::<Search>()
            .next_back()
            .expect("the user must perform a search before checking results");
        let rank = search
            .results
            .clone()
            .into_iter()
            .enumerate()
            .find_map(|(num, doc)| {
                let admin: Admin = serde_json::from_value(doc).expect("admin");
                if let Some(name) = &self.name {
                    if !equal_ignore_case_utf8(name, admin.name.as_str()) {
                        return None;
                    }
                }
                if let Some(zone_type) = &self.zone_type {
                    if !equal_ignore_case_utf8(
                        zone_type,
                        admin.zone_type.expect("zone type").as_str(),
                    ) {
                        // is it justified to expect a zone type?
                        return None;
                    }
                }
                Some(num)
            })
            .expect("document was not found in search results");

        assert!(rank < self.limit);
        Ok(StepStatus::Done)
    }
}

// Find address

#[then(
    regex = r#"he finds address "(.*)", "(.*)", "(.*)", and "(.*)" in the first "(.*)" results$"#
)]
async fn find_address(
    state: &mut GlobalState,
    house: String,
    street: String,
    city: String,
    postcode: String,
    limit: usize,
) {
    state
        .execute(HasAddress {
            house: none_if_empty(house),
            street: none_if_empty(street),
            city: none_if_empty(city),
            postcode: none_if_empty(postcode),
            limit,
        })
        .await
        .expect("failed to find document");
}

/// Check if given address is in the output.
///
/// It assumes that a Search has already been performed before.
pub struct HasAddress {
    house: Option<String>,
    street: Option<String>,
    city: Option<String>,
    postcode: Option<String>,
    limit: usize,
}

#[async_trait(?Send)]
impl Step for HasAddress {
    async fn execute(&mut self, state: &State) -> Result<StepStatus, Error> {
        let (search, _) = state
            .steps_for::<Search>()
            .next_back()
            .expect("the user must perform a search before checking results");

        let rank = search
            .results
            .clone()
            .into_iter()
            .enumerate()
            .find_map(|(num, doc)| {
                let addr: Addr = serde_json::from_value(doc).expect("address");
                if let Some(house) = &self.house {
                    if !equal_ignore_case_utf8(house, addr.house_number.as_str()) {
                        return None;
                    }
                }
                if let Some(street) = &self.street {
                    if !equal_ignore_case_utf8(street, addr.street.name.as_str()) {
                        return None;
                    }
                }
                if let Some(postcode) = &self.postcode {
                    if !equal_ignore_case_utf8(postcode, addr.zip_codes[0].as_str()) {
                        return None;
                    }
                }
                if let Some(city) = &self.city {
                    if !equal_ignore_case_utf8(city, addr.city().expect("city").as_str()) {
                        return None;
                    }
                }
                Some(num)
            })
            .expect("document was not found in search results");

        assert!(rank < self.limit);
        Ok(StepStatus::Done)
    }
}

// Find POI

#[then(
    regex = r#"he finds poi "(.*)", a "(.*)" located near ([0-9\.]+), ([0-9\.]+) in the first (\d+) results$"#
)]
async fn find_poi(
    state: &mut GlobalState,
    label: String,
    poi_type: String,
    lat: f64,
    lon: f64,
    limit: usize,
) {
    state
        .execute(HasPoi {
            label: none_if_empty(label),
            poi_type: none_if_empty(poi_type),
            coord: Some(geo_types::Point::new(lon, lat)),
            limit,
        })
        .await
        .expect("failed to find document");
}

/// Check if given poi is in the output.
///
/// It assumes that a Search has already been performed before.
pub struct HasPoi {
    label: Option<String>,
    poi_type: Option<String>,
    coord: Option<geo_types::Point<f64>>,
    limit: usize,
}

#[async_trait(?Send)]
impl Step for HasPoi {
    async fn execute(&mut self, state: &State) -> Result<StepStatus, Error> {
        let (search, _) = state
            .steps_for::<Search>()
            .next_back()
            .expect("the user must perform a search before checking results");

        let rank = search
            .results
            .clone()
            .into_iter()
            .enumerate()
            .find_map(|(num, doc)| {
                let poi: Poi = serde_json::from_value(doc).expect("poi");
                if let Some(label) = &self.label {
                    if !equal_ignore_case_utf8(label, poi.label.as_str()) {
                        return None;
                    }
                }
                if let Some(poi_type) = &self.poi_type {
                    if !equal_ignore_case_utf8(poi_type, poi.poi_type.name.as_str()) {
                        return None;
                    }
                }
                if let Some(coord) = &self.coord {
                    let distance = coord.haversine_distance(&poi.coord.into());
                    // if we are at more than 100m
                    if distance > 100.0 {
                        return None;
                    }
                }
                Some(num)
            })
            .expect("document was not found in search results");

        assert!(rank < self.limit);
        Ok(StepStatus::Done)
    }
}

// Utils

fn none_if_empty(val: String) -> Option<String> {
    if val.is_empty() {
        None
    } else {
        Some(val)
    }
}

// Based on https://stackoverflow.com/questions/63871601/what-is-an-efficient-way-to-compare-strings-while-ignoring-case
// Its a case insensitive comparison, returning true if equal, false otherwise
fn equal_ignore_case_utf8(a: &str, b: &str) -> bool {
    let abs = a
        .chars()
        .flat_map(char::to_lowercase)
        .zip_longest(b.chars().flat_map(char::to_lowercase));

    for ab in abs {
        match ab {
            Left(_) => return false,
            Right(_) => return false,
            Both(a, b) => {
                if a.cmp(&b) != Ordering::Equal {
                    return false;
                }
            }
        }
    }
    true
}
<commit_msg>fix ci<commit_after>use async_trait::async_trait;
use cucumber::{then, when};
use geo::algorithm::haversine_distance::HaversineDistance;
use itertools::{EitherOrBoth::*, Itertools};
use mimir::adapters::secondary::elasticsearch::remote::connection_test_pool;
use mimir::domain::model::configuration;
use mimir::domain::ports::secondary::remote::Remote;
use std::cmp::Ordering;

use crate::error::Error;
use crate::state::{GlobalState, State, Step, StepStatus};
use mimir::adapters::primary::bragi::api::DEFAULT_LIMIT_RESULT_ES;
use mimir::adapters::primary::{
    common::coord::Coord, common::dsl::build_query, common::filters::Filters,
    common::settings::QuerySettings,
};
use mimir::adapters::secondary::elasticsearch::ElasticsearchStorageConfig;
use mimir::domain::{model::query::Query, ports::primary::search_documents::SearchDocuments};
use places::{addr::Addr, admin::Admin, poi::Poi};

// Search place

#[when(regex = r#"the user searches (.+) datatype for "(.*)" at (.+), (.+)$"#)]
async fn search(state: &mut GlobalState, places: String, query: String, lat: f32, lon: f32) {
    perform_search(state, places, query, Coord::new(lat, lon).into(), None).await;
}

#[when(regex = r#"the user searches (.+) datatype for "(.+)" with "(.+)" filters$"#)]
async fn search_with_zone_filters(
    state: &mut GlobalState,
    places: String,
    query: String,
    zone_types: String,
) {
    perform_search(state, places, query, None, Some(zone_types)).await;
}

#[when(regex = r#"the user searches (.+) datatype for "(.*)"$"#)]
async fn search_no_coord(state: &mut GlobalState, places: String, query: String) {
    perform_search(state, places, query, None, None).await;
}

async fn perform_search(
    state: &mut GlobalState,
    places: String,
    query: String,
    coord: Option<Coord>,
    zone_types: Option<String>,
) {
    let places = {
        if places == "all" {
            vec![configuration::root()]
        } else {
            places
                .split(',')
                .map(str::trim)
                .map(str::to_string)
                .map(|s| configuration::root_doctype(&s))
                .collect()
        }
    };

    let zone_types = zone_types.map(|f| {
        f.split(',')
            .map(str::trim)
            .map(str::to_string)
            .collect::<Vec<String>>()
    });

    let filters = Filters {
        coord,
        zone_types,
        ..Default::default()
    };

    state
        .execute(Search {
            places,
            query,
            filters,
            results: Vec::new(),
        })
        .await
        .expect("failed to search");
}

/// Perform a search in current Elasticsearch DB.
#[derive(Debug)]
pub struct Search {
    places: Vec<String>, // What kind of places we're searching
    query: String,       // The search string
    filters: Filters,    // Search filters
    results: Vec<serde_json::Value>,
}

#[async_trait(?Send)]
impl Step for Search {
    async fn execute(&mut self, _state: &State) -> Result<StepStatus, Error> {
        let client = connection_test_pool()
            .conn(ElasticsearchStorageConfig::default_testing())
            .await
            .expect("Could not establish connection to Elasticsearch");

        // Build ES query
        let dsl = build_query(
            &self.query,
            self.filters.clone(),
            "fr".to_string(),
            &QuerySettings::default(),
        );

        // Fetch documents
        self.results = {
            client
                .search_documents(
                    self.places.clone(),
                    Query::QueryDSL(dsl),
                    DEFAULT_LIMIT_RESULT_ES,
                    None,
                )
                .await
                .unwrap()
        };

        Ok(StepStatus::Done)
    }
}

// Find document

#[then(regex = r#"he finds "(.+)" as the first result$"#)]
async fn find_id(state: &mut GlobalState, id: String) {
    state
        .execute(HasDocument { id, max_rank: 1 })
        .await
        .expect("failed to find document");
}

/// Check if given document is in the output.
///
/// It assumes that a Search has already been performed before.
pub struct HasDocument {
    id: String,
    max_rank: usize,
}

#[async_trait(?Send)]
impl Step for HasDocument {
    async fn execute(&mut self, state: &State) -> Result<StepStatus, Error> {
        let (search, _) = state
            .steps_for::<Search>()
            .next_back()
            .expect("the user must perform a search before checking results");

        let (rank, _) = (search.results.iter().enumerate())
            .find(|(_, doc)| {
                println!("{} vs {}", self.id, doc["id"].as_str().unwrap());
                self.id == doc["id"].as_str().unwrap()
            })
            .expect("document was not found in search results");

        assert!(rank < self.max_rank);
        Ok(StepStatus::Done)
    }
}

// Find Admin

#[then(regex = r#"he finds admin "(.*)", a "(.*)", in the first (.*) results$"#)]
async fn find_admin(state: &mut GlobalState, name: String, zone_type: String, limit: usize) {
    state
        .execute(HasAdmin {
            name: Some(name),
            zone_type: Some(zone_type),
            limit,
        })
        .await
        .expect("failed to find document");
}

/// Check if given admin is in the output.
///
/// It assumes that a Search has already been performed before.
pub struct HasAdmin {
    name: Option<String>,
    zone_type: Option<String>,
    limit: usize,
}

#[async_trait(?Send)]
impl Step for HasAdmin {
    async fn execute(&mut self, state: &State) -> Result<StepStatus, Error> {
        let (search, _) = state
            .steps_for::<Search>()
            .next_back()
            .expect("the user must perform a search before checking results");
        let rank = search
            .results
            .clone()
            .into_iter()
            .enumerate()
            .find_map(|(num, doc)| {
                let admin: Admin = serde_json::from_value(doc).expect("admin");
                if let Some(name) = &self.name {
                    if !equal_ignore_case_utf8(name, admin.name.as_str()) {
                        return None;
                    }
                }
                if let Some(zone_type) = &self.zone_type {
                    if !equal_ignore_case_utf8(
                        zone_type,
                        admin.zone_type.expect("zone type").as_str(),
                    ) {
                        // is it justified to expect a zone type?
                        return None;
                    }
                }
                Some(num)
            })
            .expect("document was not found in search results");

        assert!(rank < self.limit);
        Ok(StepStatus::Done)
    }
}

// Find address

#[then(
    regex = r#"he finds address "(.*)", "(.*)", "(.*)", and "(.*)" in the first "(.*)" results$"#
)]
async fn find_address(
    state: &mut GlobalState,
    house: String,
    street: String,
    city: String,
    postcode: String,
    limit: usize,
) {
    state
        .execute(HasAddress {
            house: none_if_empty(house),
            street: none_if_empty(street),
            city: none_if_empty(city),
            postcode: none_if_empty(postcode),
            limit,
        })
        .await
        .expect("failed to find document");
}

/// Check if given address is in the output.
///
/// It assumes that a Search has already been performed before.
pub struct HasAddress {
    house: Option<String>,
    street: Option<String>,
    city: Option<String>,
    postcode: Option<String>,
    limit: usize,
}

#[async_trait(?Send)]
impl Step for HasAddress {
    async fn execute(&mut self, state: &State) -> Result<StepStatus, Error> {
        let (search, _) = state
            .steps_for::<Search>()
            .next_back()
            .expect("the user must perform a search before checking results");

        let rank = search
            .results
            .clone()
            .into_iter()
            .enumerate()
            .find_map(|(num, doc)| {
                let addr: Addr = serde_json::from_value(doc).expect("address");
                if let Some(house) = &self.house {
                    if !equal_ignore_case_utf8(house, addr.house_number.as_str()) {
                        return None;
                    }
                }
                if let Some(street) = &self.street {
                    if !equal_ignore_case_utf8(street, addr.street.name.as_str()) {
                        return None;
                    }
                }
                if let Some(postcode) = &self.postcode {
                    if !equal_ignore_case_utf8(postcode, addr.zip_codes[0].as_str()) {
                        return None;
                    }
                }
                if let Some(city) = &self.city {
                    if !equal_ignore_case_utf8(city, addr.city().expect("city").as_str()) {
                        return None;
                    }
                }
                Some(num)
            })
            .expect("document was not found in search results");

        assert!(rank < self.limit);
        Ok(StepStatus::Done)
    }
}

// Find POI

#[then(
    regex = r#"he finds poi "(.*)", a "(.*)" located near ([0-9\.]+), ([0-9\.]+) in the first (\d+) results$"#
)]
async fn find_poi(
    state: &mut GlobalState,
    label: String,
    poi_type: String,
    lat: f64,
    lon: f64,
    limit: usize,
) {
    state
        .execute(HasPoi {
            label: none_if_empty(label),
            poi_type: none_if_empty(poi_type),
            coord: Some(geo_types::Point::new(lon, lat)),
            limit,
        })
        .await
        .expect("failed to find document");
}

/// Check if given poi is in the output.
///
/// It assumes that a Search has already been performed before.
pub struct HasPoi {
    label: Option<String>,
    poi_type: Option<String>,
    coord: Option<geo_types::Point<f64>>,
    limit: usize,
}

#[async_trait(?Send)]
impl Step for HasPoi {
    async fn execute(&mut self, state: &State) -> Result<StepStatus, Error> {
        let (search, _) = state
            .steps_for::<Search>()
            .next_back()
            .expect("the user must perform a search before checking results");

        let rank = search
            .results
            .clone()
            .into_iter()
            .enumerate()
            .find_map(|(num, doc)| {
                let poi: Poi = serde_json::from_value(doc).expect("poi");
                if let Some(label) = &self.label {
                    if !equal_ignore_case_utf8(label, poi.label.as_str()) {
                        return None;
                    }
                }
                if let Some(poi_type) = &self.poi_type {
                    if !equal_ignore_case_utf8(poi_type, poi.poi_type.name.as_str()) {
                        return None;
                    }
                }
                if let Some(coord) = &self.coord {
                    let distance = coord.haversine_distance(&poi.coord.into());
                    // if we are at more than 100m
                    if distance > 100.0 {
                        return None;
                    }
                }
                Some(num)
            })
            .expect("document was not found in search results");

        assert!(rank < self.limit);
        Ok(StepStatus::Done)
    }
}

// Utils

fn none_if_empty(val: String) -> Option<String> {
    if val.is_empty() {
        None
    } else {
        Some(val)
    }
}

// Based on https://stackoverflow.com/questions/63871601/what-is-an-efficient-way-to-compare-strings-while-ignoring-case
// Its a case insensitive comparison, returning true if equal, false otherwise
fn equal_ignore_case_utf8(a: &str, b: &str) -> bool {
    let abs = a
        .chars()
        .flat_map(char::to_lowercase)
        .zip_longest(b.chars().flat_map(char::to_lowercase));

    for ab in abs {
        match ab {
            Left(_) => return false,
            Right(_) => return false,
            Both(a, b) => {
                if a.cmp(&b) != Ordering::Equal {
                    return false;
                }
            }
        }
    }
    true
}
<|endoftext|>