<commit_before>use mozprofile::preferences::Pref;

lazy_static! {
    pub static ref DEFAULT: [(&'static str, Pref); 74] = [
        // Disable automatic downloading of new releases
        ("app.update.auto", Pref::new(false)),

        // Disable automatically upgrading Firefox
        ("app.update.enabled", Pref::new(false)),

        // Increase the APZ content response timeout in tests to 1
        // minute.  This is to accommodate the fact that test environments
        // tends to be slower than production environments (with the
        // b2g emulator being the slowest of them all), resulting in the
        // production timeout value sometimes being exceeded and causing
        // false-positive test failures.
        //
        // (bug 1176798, bug 1177018, bug 1210465)
        ("apz.content_response_timeout", Pref::new(60000)),

        // Enable the dump function, which sends messages to the system
        // console
        ("browser.dom.window.dump.enabled", Pref::new(true)),

        // Indicate that the download panel has been shown once so
        // that whichever download test runs first does not show the popup
        // inconsistently
        ("browser.download.panel.shown", Pref::new(true)),

        // Implicitly accept license
        ("browser.EULA.override", Pref::new(true)),

        // use about:blank as new tab page
        ("browser.newtabpage.enabled", Pref::new(false)),

        // Assume the about:newtab pages intro panels have been shown
        // to not depend on which test runs first and happens to open
        // about:newtab
        ("browser.newtabpage.introShown", Pref::new(true)),

        // Never start the browser in offline mode
        ("browser.offline", Pref::new(false)),

        // Background thumbnails in particular cause grief, and disabling
        // thumbnails in general cannot hurt
        ("browser.pagethumbnails.capturing_disabled", Pref::new(true)),

        // Avoid performing Reader Mode intros during tests
        ("browser.reader.detectedFirstArticle", Pref::new(true)),

        // Disable safebrowsing components
        ("browser.safebrowsing.blockedURIs.enabled", Pref::new(false)),
        ("browser.safebrowsing.downloads.enabled", Pref::new(false)),
        ("browser.safebrowsing.enabled", Pref::new(false)),
        ("browser.safebrowsing.forbiddenURIs.enabled", Pref::new(false)),
        ("browser.safebrowsing.malware.enabled", Pref::new(false)),
        ("browser.safebrowsing.phishing.enabled", Pref::new(false)),

        // Disable updates to search engines
        ("browser.search.update", Pref::new(false)),

        // Do not restore the last open set of tabs if the browser crashed
        ("browser.sessionstore.resume_from_crash", Pref::new(false)),

        // Skip check for default browser on startup
        ("browser.shell.checkDefaultBrowser", Pref::new(false)),

        // Disable Android snippets
        ("browser.snippets.enabled", Pref::new(false)),
        ("browser.snippets.syncPromo.enabled", Pref::new(false)),
        ("browser.snippets.firstrunHomepage.enabled", Pref::new(false)),

        // Do not redirect user when a milestone upgrade of Firefox
        // is detected
        ("browser.startup.homepage_override.mstone", Pref::new("ignore")),

        // Start with a blank page (about:blank)
        ("browser.startup.page", Pref::new(0)),

        // Disable tab animation
        ("browser.tabs.animate", Pref::new(false)),

        // Do not warn when quitting a window with multiple tabs
        ("browser.tabs.closeWindowWithLastTab", Pref::new(false)),

        // Do not allow background tabs to be zombified, otherwise for
        // tests that open additional tabs, the test harness tab itself
        // might get unloaded
        ("browser.tabs.disableBackgroundZombification", Pref::new(false)),

        // Do not warn when closing all other open tabs
        ("browser.tabs.warnOnCloseOtherTabs", Pref::new(false)),

        // Do not warn when multiple tabs will be opened
        ("browser.tabs.warnOnOpen", Pref::new(false)),

        // Disable first run splash page on Windows 10
        ("browser.usedOnWindows10.introURL", Pref::new("")),

        // Disable the UI tour
        ("browser.uitour.enabled", Pref::new(false)),

        // Do not show datareporting policy notifications which can
        // interfere with tests
        ("datareporting.healthreport.about.reportUrl", Pref::new("http://%(server)s/dummy/abouthealthreport/")),
        ("datareporting.healthreport.documentServerURI", Pref::new("http://%(server)s/dummy/healthreport/")),
        ("datareporting.healthreport.logging.consoleEnabled", Pref::new(false)),
        ("datareporting.healthreport.service.enabled", Pref::new(false)),
        ("datareporting.healthreport.service.firstRun", Pref::new(false)),
        ("datareporting.healthreport.uploadEnabled", Pref::new(false)),
        ("datareporting.policy.dataSubmissionEnabled", Pref::new(false)),
        ("datareporting.policy.dataSubmissionPolicyAccepted", Pref::new(false)),
        ("datareporting.policy.dataSubmissionPolicyBypassNotification", Pref::new(true)),

        // Disable popup-blocker
        ("dom.disable_open_during_load", Pref::new(false)),

        // Disable the ProcessHangMonitor
        ("dom.ipc.reportProcessHangs", Pref::new(false)),

        // Disable slow script dialogues
        ("dom.max_chrome_script_run_time", Pref::new(0)),
        ("dom.max_script_run_time", Pref::new(0)),

        // Only load extensions from the application and user profile
        // AddonManager.SCOPE_PROFILE + AddonManager.SCOPE_APPLICATION
        ("extensions.autoDisableScopes", Pref::new(0)),
        ("extensions.enabledScopes", Pref::new(5)),

        // don't block add-ons for e10s
        ("extensions.e10sBlocksEnabling", Pref::new(false)),

        // Disable metadata caching for installed add-ons by default
        ("extensions.getAddons.cache.enabled", Pref::new(false)),

        // Disable intalling any distribution extensions or add-ons
        ("extensions.installDistroAddons", Pref::new(false)),
        ("extensions.showMismatchUI", Pref::new(false)),

        // Turn off extension updates so they do not bother tests
        ("extensions.update.enabled", Pref::new(false)),
        ("extensions.update.notifyUser", Pref::new(false)),

        // Make sure opening about:addons will not hit the network
        ("extensions.webservice.discoverURL", Pref::new("http://%(server)s/dummy/discoveryURL")),

        // Allow the application to have focus even it runs in the
        // background
        ("focusmanager.testmode", Pref::new(true)),

        // Disable useragent updates
        ("general.useragent.updates.enabled", Pref::new(false)),

        // Always use network provider for geolocation tests so we bypass
        // the macOS dialog raised by the corelocation provider
        ("geo.provider.testing", Pref::new(true)),

        // Do not scan wi-fi
        ("geo.wifi.scan", Pref::new(false)),

        // No hang monitor
        ("hangmonitor.timeout", Pref::new(0)),

        // Show chrome errors and warnings in the error console
        ("javascript.options.showInConsole", Pref::new(true)),

        // Make sure the disk cache does not get auto disabled
        ("network.http.bypass-cachelock-threshold", Pref::new(200000)),

        // Do not prompt for temporary redirects
        ("network.http.prompt-temp-redirect", Pref::new(false)),

        // Disable speculative connections so they are not reported as
        // leaking when they are hanging around
        ("network.http.speculative-parallel-limit", Pref::new(0)),

        // Do not automatically switch between offline and online
        ("network.manage-offline-status", Pref::new(false)),

        // Make sure SNTP requests do not hit the network
        ("network.sntp.pools", Pref::new("%(server)s")),

        // Disable Flash.  The plugin container it is run in is
        // causing problems when quitting Firefox from geckodriver,
        // c.f. https://github.com/mozilla/geckodriver/issues/225.
        ("plugin.state.flash", Pref::new(0)),

        // Local documents have access to all other local docments,
        // including directory listings.
        ("security.fileuri.strict_origin_policy", Pref::new(false)),

        // Tests don't wait for the notification button security delay
        ("security.notification_enable_delay", Pref::new(0)),

        // Ensure blocklist updates don't hit the network
        ("services.settings.server", Pref::new("http://%(server)s/dummy/blocklist/")),

        // Do not automatically fill sign-in forms with known usernames
        // and passwords
        ("signon.autofillForms", Pref::new(false)),

        // Disable password capture, so that tests that include forms
        // are not influenced by the presence of the persistent doorhanger
        // notification
        ("signon.rememberSignons", Pref::new(false)),

        // Disable first run pages
        ("startup.homepage_welcome_url", Pref::new("about:blank")),
        ("startup.homepage_welcome_url.additional", Pref::new("")),

        // Prevent starting into safe mode after application crashes
        ("toolkit.startup.max_resumed_crashes", Pref::new(-1)),

        // We want to collect telemetry, but we don't want to send in the results
        ("toolkit.telemetry.server", Pref::new("https://%(server)s/dummy/telemetry/")),
    ];

    pub static ref REQUIRED: [(&'static str, Pref); 5] = [
        // Do not warn on quitting Firefox
        ("browser.warnOnQuit", Pref::new(false)),

        // Do not warn on exit when multiple tabs are open
        ("browser.tabs.warnOnClose", Pref::new(false)),

        // Do not warn when quitting with multiple tabs
        ("browser.showQuitWarning", Pref::new(false)),

        // Until bug 1238095 is fixed, we have to disable safe CPOW checks
        ("dom.ipc.cpows.forbid-unsafe-from-browser", Pref::new(false)),

        // TODO(ato): Should not be needed, as Marionette is enabled by
        // passing the --marionette flag to the binary
        ("marionette.defaultPrefs.enabled", Pref::new(true)),
    ];
}
<commit_msg>prefs: correct fixed array length<commit_after>use mozprofile::preferences::Pref;

lazy_static! {
    pub static ref DEFAULT: [(&'static str, Pref); 75] = [
        // Disable automatic downloading of new releases
        ("app.update.auto", Pref::new(false)),

        // Disable automatically upgrading Firefox
        ("app.update.enabled", Pref::new(false)),

        // Increase the APZ content response timeout in tests to 1
        // minute.  This is to accommodate the fact that test environments
        // tends to be slower than production environments (with the
        // b2g emulator being the slowest of them all), resulting in the
        // production timeout value sometimes being exceeded and causing
        // false-positive test failures.
        //
        // (bug 1176798, bug 1177018, bug 1210465)
        ("apz.content_response_timeout", Pref::new(60000)),

        // Enable the dump function, which sends messages to the system
        // console
        ("browser.dom.window.dump.enabled", Pref::new(true)),

        // Indicate that the download panel has been shown once so
        // that whichever download test runs first does not show the popup
        // inconsistently
        ("browser.download.panel.shown", Pref::new(true)),

        // Implicitly accept license
        ("browser.EULA.override", Pref::new(true)),

        // use about:blank as new tab page
        ("browser.newtabpage.enabled", Pref::new(false)),

        // Assume the about:newtab pages intro panels have been shown
        // to not depend on which test runs first and happens to open
        // about:newtab
        ("browser.newtabpage.introShown", Pref::new(true)),

        // Never start the browser in offline mode
        ("browser.offline", Pref::new(false)),

        // Background thumbnails in particular cause grief, and disabling
        // thumbnails in general cannot hurt
        ("browser.pagethumbnails.capturing_disabled", Pref::new(true)),

        // Avoid performing Reader Mode intros during tests
        ("browser.reader.detectedFirstArticle", Pref::new(true)),

        // Disable safebrowsing components
        ("browser.safebrowsing.blockedURIs.enabled", Pref::new(false)),
        ("browser.safebrowsing.downloads.enabled", Pref::new(false)),
        ("browser.safebrowsing.enabled", Pref::new(false)),
        ("browser.safebrowsing.forbiddenURIs.enabled", Pref::new(false)),
        ("browser.safebrowsing.malware.enabled", Pref::new(false)),
        ("browser.safebrowsing.phishing.enabled", Pref::new(false)),

        // Disable updates to search engines
        ("browser.search.update", Pref::new(false)),

        // Do not restore the last open set of tabs if the browser crashed
        ("browser.sessionstore.resume_from_crash", Pref::new(false)),

        // Skip check for default browser on startup
        ("browser.shell.checkDefaultBrowser", Pref::new(false)),

        // Disable Android snippets
        ("browser.snippets.enabled", Pref::new(false)),
        ("browser.snippets.syncPromo.enabled", Pref::new(false)),
        ("browser.snippets.firstrunHomepage.enabled", Pref::new(false)),

        // Do not redirect user when a milestone upgrade of Firefox
        // is detected
        ("browser.startup.homepage_override.mstone", Pref::new("ignore")),

        // Start with a blank page (about:blank)
        ("browser.startup.page", Pref::new(0)),

        // Disable tab animation
        ("browser.tabs.animate", Pref::new(false)),

        // Do not warn when quitting a window with multiple tabs
        ("browser.tabs.closeWindowWithLastTab", Pref::new(false)),

        // Do not allow background tabs to be zombified, otherwise for
        // tests that open additional tabs, the test harness tab itself
        // might get unloaded
        ("browser.tabs.disableBackgroundZombification", Pref::new(false)),

        // Do not warn when closing all other open tabs
        ("browser.tabs.warnOnCloseOtherTabs", Pref::new(false)),

        // Do not warn when multiple tabs will be opened
        ("browser.tabs.warnOnOpen", Pref::new(false)),

        // Disable first run splash page on Windows 10
        ("browser.usedOnWindows10.introURL", Pref::new("")),

        // Disable the UI tour
        ("browser.uitour.enabled", Pref::new(false)),

        // Do not show datareporting policy notifications which can
        // interfere with tests
        ("datareporting.healthreport.about.reportUrl", Pref::new("http://%(server)s/dummy/abouthealthreport/")),
        ("datareporting.healthreport.documentServerURI", Pref::new("http://%(server)s/dummy/healthreport/")),
        ("datareporting.healthreport.logging.consoleEnabled", Pref::new(false)),
        ("datareporting.healthreport.service.enabled", Pref::new(false)),
        ("datareporting.healthreport.service.firstRun", Pref::new(false)),
        ("datareporting.healthreport.uploadEnabled", Pref::new(false)),
        ("datareporting.policy.dataSubmissionEnabled", Pref::new(false)),
        ("datareporting.policy.dataSubmissionPolicyAccepted", Pref::new(false)),
        ("datareporting.policy.dataSubmissionPolicyBypassNotification", Pref::new(true)),

        // Disable popup-blocker
        ("dom.disable_open_during_load", Pref::new(false)),

        // Disable the ProcessHangMonitor
        ("dom.ipc.reportProcessHangs", Pref::new(false)),

        // Disable slow script dialogues
        ("dom.max_chrome_script_run_time", Pref::new(0)),
        ("dom.max_script_run_time", Pref::new(0)),

        // Only load extensions from the application and user profile
        // AddonManager.SCOPE_PROFILE + AddonManager.SCOPE_APPLICATION
        ("extensions.autoDisableScopes", Pref::new(0)),
        ("extensions.enabledScopes", Pref::new(5)),

        // don't block add-ons for e10s
        ("extensions.e10sBlocksEnabling", Pref::new(false)),

        // Disable metadata caching for installed add-ons by default
        ("extensions.getAddons.cache.enabled", Pref::new(false)),

        // Disable intalling any distribution extensions or add-ons
        ("extensions.installDistroAddons", Pref::new(false)),
        ("extensions.showMismatchUI", Pref::new(false)),

        // Turn off extension updates so they do not bother tests
        ("extensions.update.enabled", Pref::new(false)),
        ("extensions.update.notifyUser", Pref::new(false)),

        // Make sure opening about:addons will not hit the network
        ("extensions.webservice.discoverURL", Pref::new("http://%(server)s/dummy/discoveryURL")),

        // Allow the application to have focus even it runs in the
        // background
        ("focusmanager.testmode", Pref::new(true)),

        // Disable useragent updates
        ("general.useragent.updates.enabled", Pref::new(false)),

        // Always use network provider for geolocation tests so we bypass
        // the macOS dialog raised by the corelocation provider
        ("geo.provider.testing", Pref::new(true)),

        // Do not scan wi-fi
        ("geo.wifi.scan", Pref::new(false)),

        // No hang monitor
        ("hangmonitor.timeout", Pref::new(0)),

        // Show chrome errors and warnings in the error console
        ("javascript.options.showInConsole", Pref::new(true)),

        // Make sure the disk cache does not get auto disabled
        ("network.http.bypass-cachelock-threshold", Pref::new(200000)),

        // Do not prompt for temporary redirects
        ("network.http.prompt-temp-redirect", Pref::new(false)),

        // Disable speculative connections so they are not reported as
        // leaking when they are hanging around
        ("network.http.speculative-parallel-limit", Pref::new(0)),

        // Do not automatically switch between offline and online
        ("network.manage-offline-status", Pref::new(false)),

        // Make sure SNTP requests do not hit the network
        ("network.sntp.pools", Pref::new("%(server)s")),

        // Disable Flash.  The plugin container it is run in is
        // causing problems when quitting Firefox from geckodriver,
        // c.f. https://github.com/mozilla/geckodriver/issues/225.
        ("plugin.state.flash", Pref::new(0)),

        // Local documents have access to all other local docments,
        // including directory listings.
        ("security.fileuri.strict_origin_policy", Pref::new(false)),

        // Tests don't wait for the notification button security delay
        ("security.notification_enable_delay", Pref::new(0)),

        // Ensure blocklist updates don't hit the network
        ("services.settings.server", Pref::new("http://%(server)s/dummy/blocklist/")),

        // Do not automatically fill sign-in forms with known usernames
        // and passwords
        ("signon.autofillForms", Pref::new(false)),

        // Disable password capture, so that tests that include forms
        // are not influenced by the presence of the persistent doorhanger
        // notification
        ("signon.rememberSignons", Pref::new(false)),

        // Disable first run pages
        ("startup.homepage_welcome_url", Pref::new("about:blank")),
        ("startup.homepage_welcome_url.additional", Pref::new("")),

        // Prevent starting into safe mode after application crashes
        ("toolkit.startup.max_resumed_crashes", Pref::new(-1)),

        // We want to collect telemetry, but we don't want to send in the results
        ("toolkit.telemetry.server", Pref::new("https://%(server)s/dummy/telemetry/")),
    ];

    pub static ref REQUIRED: [(&'static str, Pref); 5] = [
        // Do not warn on quitting Firefox
        ("browser.warnOnQuit", Pref::new(false)),

        // Do not warn on exit when multiple tabs are open
        ("browser.tabs.warnOnClose", Pref::new(false)),

        // Do not warn when quitting with multiple tabs
        ("browser.showQuitWarning", Pref::new(false)),

        // Until bug 1238095 is fixed, we have to disable safe CPOW checks
        ("dom.ipc.cpows.forbid-unsafe-from-browser", Pref::new(false)),

        // TODO(ato): Should not be needed, as Marionette is enabled by
        // passing the --marionette flag to the binary
        ("marionette.defaultPrefs.enabled", Pref::new(true)),
    ];
}
<|endoftext|>


<commit_before>use std::io;
use std::net::SocketAddr;
use std::str;
use std::time::Duration;

use futures::{future, Future, Stream};
use tokio_core::reactor::{Core, Handle};
use tokio_core::net::{TcpListener, TcpStream};
use tokio_timer::Timer;
use hyper::{self, Headers, Client, HttpVersion};
use hyper::client;
use hyper::client::Service;
use hyper::header;
use hyper::server::{self, Http};
use hyper_tls::HttpsConnector;
use hyper::{Url, Uri};

use pool::Pool;
use mgmt::Mgmt;
use stream::{merge3, Merged3Item};
use health;

// testing here before sending PR upstream
// TODO make this typed
header! { (Via, "Via") => [String] }
header! { (TE, "TE") => [String] }
header! { (ProxyAuthorization, "Proxy-Authorization") => [String] }
header! { (ProxyAuthenticate, "Proxy-Authenticate") => [String] }
header! { (Trailer, "Trailer") => [String] }

impl Via {

    /// Append a Via header to existing Via header
    ///
    /// This is used when the upstream sends a Via header and we need to create a list of Via
    /// header values.
    pub fn append(&mut self, other: Via) {
        let s = format!(", {}", other);
        self.0.push_str(&s);
    }
}

/// Create Via header for proxy to send downstream to origin server
///
/// The Via header may already exist, so create a new header based off the upstream value
pub fn create_via_header(via: Option<& Via>, version: &HttpVersion) -> Via {

    let version = match version {
        &HttpVersion::Http09 => "0.9",
        &HttpVersion::Http10 => "1.0",
        &HttpVersion::Http11 => "1.1",
        &HttpVersion::H2 => "2",
        &HttpVersion::H2c => "2",
        _ => unreachable!(),
    };

    let value = Via(format!("{} weldr", version));

    match via {
        Some(v) => {
            let mut v = v.clone();
            v.append(value);
            v
        }
        None => {
            value
        }
    }
}

/// Remove frontend request headers that should not be sent to backend
///
/// This creates a new collection rather than modify the existing one. Hyper can sometimes return
/// `None` even though a header is removed. Thus, we ignore the result.
pub fn filter_frontend_request_headers(headers: &Headers) -> Headers {

    let mut h = headers.clone();

    headers.get::<header::Connection>().and_then(|c| {
        for c_h in &c.0 {
            match c_h {
                &header::ConnectionOption::Close => {
                    let _ = h.remove_raw("Close");
                }

                &header::ConnectionOption::KeepAlive => {
                    let _ = h.remove_raw("Keep-Alive");
                }

                &header::ConnectionOption::ConnectionHeader(ref o) => {
                    let _ = h.remove_raw(&o);
                }
            }
        }

        Some(())
    });

    let _ = h.remove::<header::Connection>();
    let _ = h.remove::<TE>();
    let _ = h.remove::<header::TransferEncoding>();
    let _ = h.remove::<ProxyAuthorization>();
    let _ = h.remove::<Trailer>();
    let _ = h.remove::<header::Upgrade>();

    h
}

fn map_uri_to_url(uri: &Uri) -> Url {
    debug!("uri = {:?}", uri);
    Url::parse(
        &format!(
            "{}://{}{}?{}",
            uri.scheme().unwrap_or("http"),
            uri.authority().unwrap_or("example.com"),
            uri.path(),
            uri.query().unwrap_or(""),
        )
    ).expect("Failed to map uri to url")
}

/// Map a frontend request to a backend request
///
/// The primary purpose of this function is to add and remove headers as required by an
/// intermediary conforming to the HTTP spec.
fn map_request(req: server::Request) -> client::Request {
    let via = create_via_header(
        req.headers().get::<Via>(),
        req.version());

    let mut headers = filter_frontend_request_headers(req.headers());
    headers.set(via);

    //if map_host {
    //    // add host header related to backend
    //    let _ = headers.remove::<header::Host>();
    //    let host = url.host_str().unwrap().to_string();
    //    let port = url.port_or_known_default();
    //    headers.set(
    //        header::Host::new(host, port)
    //    );
    //}

    let url = map_uri_to_url(req.uri());

    let mut r = client::Request::new(req.method().clone(), url);
    r.headers_mut().extend(headers.iter());
    r.set_body(req.body());
    r
}

pub fn filter_backend_response_headers(headers: &Headers) -> Headers {
    let mut h = headers.clone();

    let _ = h.remove::<header::TransferEncoding>();
    let _ = h.remove::<ProxyAuthenticate>();
    let _ = h.remove::<Trailer>();
    let _ = h.remove::<header::Upgrade>();

    h
}

/// Map a backend response to a frontend response
///
/// The primary purpose of this function is to add and remove headers as required by an
/// intermediary conforming to the HTTP spec.
fn map_response(res: client::Response) -> server::Response {
    let mut r = server::Response::new().with_status(*res.status());

    let headers = filter_backend_response_headers(res.headers());
    r.headers_mut().extend(headers.iter());

    r.set_body(res.body());
    r
}

struct Proxy {
    client: Client<HttpsConnector>,
    pool: Pool,
}

impl Service for Proxy {
    type Request = server::Request;
    type Response = server::Response;
    type Error = hyper::Error;
    type Future = Box<Future<Item=server::Response, Error=Self::Error>>;

    fn call(&self, req: server::Request) -> Self::Future {

        let mut client_req = map_request(req);

        //let server = self.pool.get().expect("Failed to get server from pool");
        self.pool.request(|server| {

            // TODO need to add query strings in here as well
            let url = server.url().join(client_req.url().path()).unwrap();
            let map_host = server.map_host();
            debug!("Preparing backend request to {:?}", url);

            if map_host {
                // add host header related to backend
                let _ = client_req.headers_mut().remove::<header::Host>();
                let host = url.host_str().unwrap().to_string();
                let port = url.port_or_known_default();
                client_req.headers_mut().set(header::Host::new(host, port));
            }
            client_req.set_url(url);

            let backend = self.client.call(client_req).then(move |res| {
                match res {
                    Ok(res) => {
                        debug!("Response: {}", res.status());
                        debug!("Headers: \n{}", res.headers());

                        let server_response = map_response(res);

                        ::futures::finished(server_response)
                    }
                    Err(e) => {
                        error!("Error connecting to backend: {:?}", e);
                        ::futures::failed(e)
                    }
                }
            });

            Box::new(backend)
        })
    }
}

/// Run server with default Core
pub fn run(proxy_addr: SocketAddr, admin_addr: SocketAddr, pool: Pool) -> io::Result<()> {
    let core = Core::new()?;
    let handle = core.handle();

    let listener = TcpListener::bind(&proxy_addr, &handle)?;
    let admin_listener = TcpListener::bind(&admin_addr, &handle)?;
    run_with(core, listener, admin_listener, pool, future::empty())
}

/// Run server with specified Core, TcpListener, Pool
///
/// This is useful for integration testing where the port is set to 0 and the test code needs to
/// determine the local addr.
pub fn run_with<F>(mut core: Core, listener: TcpListener, admin_listener: TcpListener, pool: Pool, shutdown_signal: F) -> io::Result<()>
    where F: Future<Item = (), Error = hyper::Error>,
{
    let handle = core.handle();

    let timer = Timer::default();

    // FIXME configure health check timer
    let health_timer = timer.interval(Duration::from_secs(5)).map_err(|e| {
        io::Error::new(io::ErrorKind::Other, e)
    });

    let local_addr = listener.local_addr()?;
    let listener = merge3(listener.incoming(), admin_listener.incoming(), health_timer);
    let srv = listener.for_each(move |stream| {

        // first stream is the proxy ip
        // second stream is the management ip
        // third stream is health interval
        match stream {
            Merged3Item::First((socket, addr)) => {
                proxy(socket, addr, pool.clone(), &handle);
            }
            Merged3Item::Second((socket, addr)) => {
                mgmt(socket, addr, pool.clone(), &handle);
            }
            Merged3Item::Third(()) => {
                info!("health check");
                health::run(pool.clone(), &handle);
            }
            Merged3Item::FirstSecond((socket, addr), (socket2, addr2)) => {
                proxy(socket, addr, pool.clone(), &handle);
                mgmt(socket2, addr2, pool.clone(), &handle);
            }
            Merged3Item::SecondThird((socket, addr), ()) => {
                mgmt(socket, addr, pool.clone(), &handle);
                info!("health check");
                health::run(pool.clone(), &handle);
            }
            Merged3Item::FirstThird((socket, addr), ()) => {
                proxy(socket, addr, pool.clone(), &handle);
                info!("health check");
                health::run(pool.clone(), &handle);
            }
            Merged3Item::All((socket, addr), (socket2, addr2), ()) => {
                proxy(socket, addr, pool.clone(), &handle);
                mgmt(socket2, addr2, pool.clone(), &handle);
                info!("health check");
                health::run(pool.clone(), &handle);
            }
        }

        Ok(())
    });

    info!("Listening on http://{}", &local_addr);
    match core.run(shutdown_signal.select(srv.map_err(|e| e.into()))) {
        Ok(((), _incoming)) => Ok(()),
        Err((e, _other)) => return Err(io::Error::new(io::ErrorKind::Other, e)),
    }
}

fn proxy(socket: TcpStream, addr: SocketAddr, pool: Pool, handle: &Handle) {
    let client = Client::configure()
        .connector(HttpsConnector::new(4, handle))
        .build(&handle);
    let service = Proxy {
        client: client,
        pool: pool,
    };

    let http = Http::new();
    http.bind_connection(&handle, socket, addr, service);
}

fn mgmt(socket: TcpStream, addr: SocketAddr, pool: Pool, handle: &Handle) {
    let service = Mgmt::new(pool, handle.clone());
    let http = Http::new();
    http.bind_connection(&handle, socket, addr, service);
}

#[cfg(test)]
mod tests {
    use hyper::{HttpVersion, Headers};
    use hyper::header;

    use super::*;

    #[test]
    #[ignore]
    /// Send HTTP 1.0 request and ensure proxy sends HTTP 1.1
    ///
    /// Per RFC 7230 Section 2.6 - MUST send own HTTP version
    fn test_must_send_own_http_version() {
        // logic is in, waiting for mocked request support to test
    }

    #[test]
    #[ignore]
    /// Send HTTP request and ensure proxy sets proper Via header
    ///
    /// Per RFC 7230 Section 5.7.1
    fn test_proxy_sets_via_header() {
        // logic is in, waiting for mocked request support to test
    }

    #[test]
    fn test_create_via_header() {
        let via = Via("1.0 proxy".to_owned());
        let version = HttpVersion::Http11;

        let given = create_via_header(None, &version);

        assert_eq!(Via("1.1 weldr".to_owned()), given);

        let given = create_via_header(Some(&via), &version);

        assert_eq!(Via("1.0 proxy, 1.1 weldr".to_owned()), given);
    }

    #[test]
    /// Per RFC 2616 Section 13.5.1 - MUST remove hop-by-hop headers
    /// Per RFC 7230 Section 6.1 - MUST remove Connection and Connection option headers
    fn test_filter_frontend_request_headers() {

        let bad = vec![
            ("TE", "gzip"),
            ("Transfer-Encoding", "chunked"),
            ("Host", "example.net"),
            ("Connection", "Keep-Alive, Foo, Bar"),
            ("Foo", "abc"),
            ("Foo", "def"),
            ("Keep-Alive", "timeout=30"),
            ("Proxy-Authorization", "randombase64value"),
            ("Trailer", "X-Random-Header"),
            ("Upgrade", "HTTP/2.0, SHTTP/1.3, IRC/6.9, RTA/x11"),
        ];

        let mut headers = Headers::new();

        for (name, value) in bad {
            headers.set_raw(name, value);
        }

        let given = filter_frontend_request_headers(&headers);

        // defining these here only to let me assert
        header! { (Foo, "Foo") => [String] }
        header! { (KeepAlive, "Keep-Alive") => [String] }

        assert_eq!(false, given.has::<TE>());
        assert_eq!(false, given.has::<header::TransferEncoding>());
        assert_eq!(true, given.has::<header::Host>());
        assert_eq!(false, given.has::<header::Connection>());
        assert_eq!(false, given.has::<Foo>());
        assert_eq!(false, given.has::<KeepAlive>());
        assert_eq!(false, given.has::<ProxyAuthorization>());
        assert_eq!(false, given.has::<Trailer>());
        assert_eq!(false, given.has::<header::Upgrade>());
    }

    #[test]
    /// Per RFC 2616 Section 13.5.1 - MUST remove hop-by-hop headers
    fn test_filter_backend_response_headers() {

        let bad = vec![
            ("Transfer-Encoding", "chunked"),
            ("Host", "example.net"),
            ("Proxy-Authenticate", "randombase64value"),
            ("Trailer", "X-Random-Header"),
            ("Upgrade", "HTTP/2.0"),
        ];

        let mut headers = Headers::new();

        for (name, value) in bad {
            headers.set_raw(name, value);
        }

        let given = filter_backend_response_headers(&headers);

        assert_eq!(false, given.has::<header::TransferEncoding>());
        assert_eq!(true, given.has::<header::Host>());
        assert_eq!(false, given.has::<ProxyAuthenticate>());
        assert_eq!(false, given.has::<Trailer>());
        assert_eq!(false, given.has::<header::Upgrade>());
    }
}
<commit_msg>Do not delay incoming TCP connections<commit_after>use std::io;
use std::net::SocketAddr;
use std::str;
use std::time::Duration;

use futures::{future, Future, Stream};
use tokio_core::reactor::{Core, Handle};
use tokio_core::net::{TcpListener, TcpStream};
use tokio_timer::Timer;
use hyper::{self, Headers, Client, HttpVersion};
use hyper::client;
use hyper::client::Service;
use hyper::header;
use hyper::server::{self, Http};
use hyper_tls::HttpsConnector;
use hyper::{Url, Uri};

use pool::Pool;
use mgmt::Mgmt;
use stream::{merge3, Merged3Item};
use health;

// testing here before sending PR upstream
// TODO make this typed
header! { (Via, "Via") => [String] }
header! { (TE, "TE") => [String] }
header! { (ProxyAuthorization, "Proxy-Authorization") => [String] }
header! { (ProxyAuthenticate, "Proxy-Authenticate") => [String] }
header! { (Trailer, "Trailer") => [String] }

impl Via {

    /// Append a Via header to existing Via header
    ///
    /// This is used when the upstream sends a Via header and we need to create a list of Via
    /// header values.
    pub fn append(&mut self, other: Via) {
        let s = format!(", {}", other);
        self.0.push_str(&s);
    }
}

/// Create Via header for proxy to send downstream to origin server
///
/// The Via header may already exist, so create a new header based off the upstream value
pub fn create_via_header(via: Option<& Via>, version: &HttpVersion) -> Via {

    let version = match version {
        &HttpVersion::Http09 => "0.9",
        &HttpVersion::Http10 => "1.0",
        &HttpVersion::Http11 => "1.1",
        &HttpVersion::H2 => "2",
        &HttpVersion::H2c => "2",
        _ => unreachable!(),
    };

    let value = Via(format!("{} weldr", version));

    match via {
        Some(v) => {
            let mut v = v.clone();
            v.append(value);
            v
        }
        None => {
            value
        }
    }
}

/// Remove frontend request headers that should not be sent to backend
///
/// This creates a new collection rather than modify the existing one. Hyper can sometimes return
/// `None` even though a header is removed. Thus, we ignore the result.
pub fn filter_frontend_request_headers(headers: &Headers) -> Headers {

    let mut h = headers.clone();

    headers.get::<header::Connection>().and_then(|c| {
        for c_h in &c.0 {
            match c_h {
                &header::ConnectionOption::Close => {
                    let _ = h.remove_raw("Close");
                }

                &header::ConnectionOption::KeepAlive => {
                    let _ = h.remove_raw("Keep-Alive");
                }

                &header::ConnectionOption::ConnectionHeader(ref o) => {
                    let _ = h.remove_raw(&o);
                }
            }
        }

        Some(())
    });

    let _ = h.remove::<header::Connection>();
    let _ = h.remove::<TE>();
    let _ = h.remove::<header::TransferEncoding>();
    let _ = h.remove::<ProxyAuthorization>();
    let _ = h.remove::<Trailer>();
    let _ = h.remove::<header::Upgrade>();

    h
}

fn map_uri_to_url(uri: &Uri) -> Url {
    debug!("uri = {:?}", uri);
    Url::parse(
        &format!(
            "{}://{}{}?{}",
            uri.scheme().unwrap_or("http"),
            uri.authority().unwrap_or("example.com"),
            uri.path(),
            uri.query().unwrap_or(""),
        )
    ).expect("Failed to map uri to url")
}

/// Map a frontend request to a backend request
///
/// The primary purpose of this function is to add and remove headers as required by an
/// intermediary conforming to the HTTP spec.
fn map_request(req: server::Request) -> client::Request {
    let via = create_via_header(
        req.headers().get::<Via>(),
        req.version());

    let mut headers = filter_frontend_request_headers(req.headers());
    headers.set(via);

    //if map_host {
    //    // add host header related to backend
    //    let _ = headers.remove::<header::Host>();
    //    let host = url.host_str().unwrap().to_string();
    //    let port = url.port_or_known_default();
    //    headers.set(
    //        header::Host::new(host, port)
    //    );
    //}

    let url = map_uri_to_url(req.uri());

    let mut r = client::Request::new(req.method().clone(), url);
    r.headers_mut().extend(headers.iter());
    r.set_body(req.body());
    r
}

pub fn filter_backend_response_headers(headers: &Headers) -> Headers {
    let mut h = headers.clone();

    let _ = h.remove::<header::TransferEncoding>();
    let _ = h.remove::<ProxyAuthenticate>();
    let _ = h.remove::<Trailer>();
    let _ = h.remove::<header::Upgrade>();

    h
}

/// Map a backend response to a frontend response
///
/// The primary purpose of this function is to add and remove headers as required by an
/// intermediary conforming to the HTTP spec.
fn map_response(res: client::Response) -> server::Response {
    let mut r = server::Response::new().with_status(*res.status());

    let headers = filter_backend_response_headers(res.headers());
    r.headers_mut().extend(headers.iter());

    r.set_body(res.body());
    r
}

struct Proxy {
    client: Client<HttpsConnector>,
    pool: Pool,
}

impl Service for Proxy {
    type Request = server::Request;
    type Response = server::Response;
    type Error = hyper::Error;
    type Future = Box<Future<Item=server::Response, Error=Self::Error>>;

    fn call(&self, req: server::Request) -> Self::Future {

        let mut client_req = map_request(req);

        //let server = self.pool.get().expect("Failed to get server from pool");
        self.pool.request(|server| {

            // TODO need to add query strings in here as well
            let url = server.url().join(client_req.url().path()).unwrap();
            let map_host = server.map_host();
            debug!("Preparing backend request to {:?}", url);

            if map_host {
                // add host header related to backend
                let _ = client_req.headers_mut().remove::<header::Host>();
                let host = url.host_str().unwrap().to_string();
                let port = url.port_or_known_default();
                client_req.headers_mut().set(header::Host::new(host, port));
            }
            client_req.set_url(url);

            let backend = self.client.call(client_req).then(move |res| {
                match res {
                    Ok(res) => {
                        debug!("Response: {}", res.status());
                        debug!("Headers: \n{}", res.headers());

                        let server_response = map_response(res);

                        ::futures::finished(server_response)
                    }
                    Err(e) => {
                        error!("Error connecting to backend: {:?}", e);
                        ::futures::failed(e)
                    }
                }
            });

            Box::new(backend)
        })
    }
}

/// Run server with default Core
pub fn run(proxy_addr: SocketAddr, admin_addr: SocketAddr, pool: Pool) -> io::Result<()> {
    let core = Core::new()?;
    let handle = core.handle();

    let listener = TcpListener::bind(&proxy_addr, &handle)?;
    let admin_listener = TcpListener::bind(&admin_addr, &handle)?;
    run_with(core, listener, admin_listener, pool, future::empty())
}

/// Run server with specified Core, TcpListener, Pool
///
/// This is useful for integration testing where the port is set to 0 and the test code needs to
/// determine the local addr.
pub fn run_with<F>(mut core: Core, listener: TcpListener, admin_listener: TcpListener, pool: Pool, shutdown_signal: F) -> io::Result<()>
    where F: Future<Item = (), Error = hyper::Error>,
{
    let handle = core.handle();

    let timer = Timer::default();

    // FIXME configure health check timer
    let health_timer = timer.interval(Duration::from_secs(5)).map_err(|e| {
        io::Error::new(io::ErrorKind::Other, e)
    });

    let local_addr = listener.local_addr()?;
    let listener = merge3(listener.incoming(), admin_listener.incoming(), health_timer);
    let srv = listener.for_each(move |stream| {

        // first stream is the proxy ip
        // second stream is the management ip
        // third stream is health interval
        match stream {
            Merged3Item::First((socket, addr)) => {
                proxy(socket, addr, pool.clone(), &handle);
            }
            Merged3Item::Second((socket, addr)) => {
                mgmt(socket, addr, pool.clone(), &handle);
            }
            Merged3Item::Third(()) => {
                info!("health check");
                health::run(pool.clone(), &handle);
            }
            Merged3Item::FirstSecond((socket, addr), (socket2, addr2)) => {
                proxy(socket, addr, pool.clone(), &handle);
                mgmt(socket2, addr2, pool.clone(), &handle);
            }
            Merged3Item::SecondThird((socket, addr), ()) => {
                mgmt(socket, addr, pool.clone(), &handle);
                info!("health check");
                health::run(pool.clone(), &handle);
            }
            Merged3Item::FirstThird((socket, addr), ()) => {
                proxy(socket, addr, pool.clone(), &handle);
                info!("health check");
                health::run(pool.clone(), &handle);
            }
            Merged3Item::All((socket, addr), (socket2, addr2), ()) => {
                proxy(socket, addr, pool.clone(), &handle);
                mgmt(socket2, addr2, pool.clone(), &handle);
                info!("health check");
                health::run(pool.clone(), &handle);
            }
        }

        Ok(())
    });

    info!("Listening on http://{}", &local_addr);
    match core.run(shutdown_signal.select(srv.map_err(|e| e.into()))) {
        Ok(((), _incoming)) => Ok(()),
        Err((e, _other)) => return Err(io::Error::new(io::ErrorKind::Other, e)),
    }
}

fn proxy(socket: TcpStream, addr: SocketAddr, pool: Pool, handle: &Handle) {

    // disable Nagle's algo
    // https://github.com/hyperium/hyper/issues/944
    socket.set_nodelay(true).unwrap();
    let client = Client::configure()
        .connector(HttpsConnector::new(4, handle))
        .build(&handle);
    let service = Proxy {
        client: client,
        pool: pool,
    };

    let http = Http::new();
    http.bind_connection(&handle, socket, addr, service);
}

fn mgmt(socket: TcpStream, addr: SocketAddr, pool: Pool, handle: &Handle) {
    let service = Mgmt::new(pool, handle.clone());
    let http = Http::new();
    http.bind_connection(&handle, socket, addr, service);
}

#[cfg(test)]
mod tests {
    use hyper::{HttpVersion, Headers};
    use hyper::header;

    use super::*;

    #[test]
    #[ignore]
    /// Send HTTP 1.0 request and ensure proxy sends HTTP 1.1
    ///
    /// Per RFC 7230 Section 2.6 - MUST send own HTTP version
    fn test_must_send_own_http_version() {
        // logic is in, waiting for mocked request support to test
    }

    #[test]
    #[ignore]
    /// Send HTTP request and ensure proxy sets proper Via header
    ///
    /// Per RFC 7230 Section 5.7.1
    fn test_proxy_sets_via_header() {
        // logic is in, waiting for mocked request support to test
    }

    #[test]
    fn test_create_via_header() {
        let via = Via("1.0 proxy".to_owned());
        let version = HttpVersion::Http11;

        let given = create_via_header(None, &version);

        assert_eq!(Via("1.1 weldr".to_owned()), given);

        let given = create_via_header(Some(&via), &version);

        assert_eq!(Via("1.0 proxy, 1.1 weldr".to_owned()), given);
    }

    #[test]
    /// Per RFC 2616 Section 13.5.1 - MUST remove hop-by-hop headers
    /// Per RFC 7230 Section 6.1 - MUST remove Connection and Connection option headers
    fn test_filter_frontend_request_headers() {

        let bad = vec![
            ("TE", "gzip"),
            ("Transfer-Encoding", "chunked"),
            ("Host", "example.net"),
            ("Connection", "Keep-Alive, Foo, Bar"),
            ("Foo", "abc"),
            ("Foo", "def"),
            ("Keep-Alive", "timeout=30"),
            ("Proxy-Authorization", "randombase64value"),
            ("Trailer", "X-Random-Header"),
            ("Upgrade", "HTTP/2.0, SHTTP/1.3, IRC/6.9, RTA/x11"),
        ];

        let mut headers = Headers::new();

        for (name, value) in bad {
            headers.set_raw(name, value);
        }

        let given = filter_frontend_request_headers(&headers);

        // defining these here only to let me assert
        header! { (Foo, "Foo") => [String] }
        header! { (KeepAlive, "Keep-Alive") => [String] }

        assert_eq!(false, given.has::<TE>());
        assert_eq!(false, given.has::<header::TransferEncoding>());
        assert_eq!(true, given.has::<header::Host>());
        assert_eq!(false, given.has::<header::Connection>());
        assert_eq!(false, given.has::<Foo>());
        assert_eq!(false, given.has::<KeepAlive>());
        assert_eq!(false, given.has::<ProxyAuthorization>());
        assert_eq!(false, given.has::<Trailer>());
        assert_eq!(false, given.has::<header::Upgrade>());
    }

    #[test]
    /// Per RFC 2616 Section 13.5.1 - MUST remove hop-by-hop headers
    fn test_filter_backend_response_headers() {

        let bad = vec![
            ("Transfer-Encoding", "chunked"),
            ("Host", "example.net"),
            ("Proxy-Authenticate", "randombase64value"),
            ("Trailer", "X-Random-Header"),
            ("Upgrade", "HTTP/2.0"),
        ];

        let mut headers = Headers::new();

        for (name, value) in bad {
            headers.set_raw(name, value);
        }

        let given = filter_backend_response_headers(&headers);

        assert_eq!(false, given.has::<header::TransferEncoding>());
        assert_eq!(true, given.has::<header::Host>());
        assert_eq!(false, given.has::<ProxyAuthenticate>());
        assert_eq!(false, given.has::<Trailer>());
        assert_eq!(false, given.has::<header::Upgrade>());
    }
}
<|endoftext|>


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

use super::Document;

#[derive(Debug, PartialEq)]
pub enum FilterParseError {
    ExpectedObject,
    ExpectedString,
    ExpectedArray,
    UnknownFilterType(String),
    NoFilter,
}

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

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

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

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

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

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

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

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

                match obj.get(field) {
                    Some(&Json::Null) => true,
                    None => true,
                    _ => false,
                }
            }
            Filter::And{ref children} => {
                for child in children.iter() {
                    if !child.matches(doc) {
                        return false;
                    }
                }

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

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

pub fn parse_filter(json: &Json) -> Result<Filter, FilterParseError> {
    let filter_json = try!(json.as_object().ok_or(FilterParseError::ExpectedObject));
    let first_key = try!(filter_json.keys().nth(0).ok_or(FilterParseError::NoFilter));

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

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

        Ok(Filter::Prefix{
            field: first_key.clone(),
            value: value.to_owned()
        })
    } else if first_key == "missing" {
        let filter_json = filter_json.get("missing").unwrap().as_object().unwrap();
        let first_key = filter_json.keys().nth(0).unwrap();

        Ok(Filter::Missing{
            field: first_key.clone()
        })
    } else if first_key == "and" {
        Ok(Filter::And{
            children: filter_json.get("and").unwrap()
                       .as_array().unwrap()
                       .iter().map(|f| parse_filter(f).unwrap())
                       .collect::<Vec<_>>()
        })
    } else if first_key == "or" {
        Ok(Filter::Or{
            children: filter_json.get("or").unwrap()
                       .as_array().unwrap()
                       .iter().map(|f| parse_filter(f).unwrap())
                       .collect::<Vec<_>>()
        })
    } else if first_key == "not" {
        Ok(Filter::Not{
            child: Box::new(parse_filter(filter_json.get("not").unwrap()).unwrap())
        })
    } else {
        Err(FilterParseError::UnknownFilterType(first_key.clone()))
    }
}

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

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

                    return field_value.contains(&query);
                }

                false
            }
            Query::MultiMatch{ref fields, ref query} => {
                let obj = doc.data.as_object().unwrap();

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

                        if field_value.contains(&query) {
                            return true;
                        }
                    }
                }

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

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

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

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

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

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

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

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

    let filter_json = try!(json_object.get("filter").ok_or(QueryParseError::FilteredNoFilter));
    let filter = match parse_filter(filter_json) {
        Ok(filter) => filter,
        Err(err) => return Err(QueryParseError::FilterParseError(err)),
    };

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

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

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

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

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

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

        assert_eq!(query, Ok(Query::MatchAll))
    }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        assert_eq!(query, Err(QueryParseError::FilteredNoFilter))
    }
}
<commit_msg>Fixed test<commit_after>use rustc_serialize::json::Json;

use super::Document;

#[derive(Debug, PartialEq)]
pub enum FilterParseError {
    ExpectedObject,
    ExpectedString,
    ExpectedArray,
    UnknownFilterType(String),
    NoFilter,
}

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

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

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

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

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

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

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

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

                match obj.get(field) {
                    Some(&Json::Null) => true,
                    None => true,
                    _ => false,
                }
            }
            Filter::And{ref children} => {
                for child in children.iter() {
                    if !child.matches(doc) {
                        return false;
                    }
                }

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

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

pub fn parse_filter(json: &Json) -> Result<Filter, FilterParseError> {
    let filter_json = try!(json.as_object().ok_or(FilterParseError::ExpectedObject));
    let first_key = try!(filter_json.keys().nth(0).ok_or(FilterParseError::NoFilter));

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

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

        Ok(Filter::Prefix{
            field: first_key.clone(),
            value: value.to_owned()
        })
    } else if first_key == "missing" {
        let filter_json = filter_json.get("missing").unwrap().as_object().unwrap();
        let first_key = filter_json.keys().nth(0).unwrap();

        Ok(Filter::Missing{
            field: first_key.clone()
        })
    } else if first_key == "and" {
        Ok(Filter::And{
            children: filter_json.get("and").unwrap()
                       .as_array().unwrap()
                       .iter().map(|f| parse_filter(f).unwrap())
                       .collect::<Vec<_>>()
        })
    } else if first_key == "or" {
        Ok(Filter::Or{
            children: filter_json.get("or").unwrap()
                       .as_array().unwrap()
                       .iter().map(|f| parse_filter(f).unwrap())
                       .collect::<Vec<_>>()
        })
    } else if first_key == "not" {
        Ok(Filter::Not{
            child: Box::new(parse_filter(filter_json.get("not").unwrap()).unwrap())
        })
    } else {
        Err(FilterParseError::UnknownFilterType(first_key.clone()))
    }
}

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

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

                    return field_value.contains(&query);
                }

                false
            }
            Query::MultiMatch{ref fields, ref query} => {
                let obj = doc.data.as_object().unwrap();

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

                        if field_value.contains(&query) {
                            return true;
                        }
                    }
                }

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

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

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

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

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

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

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

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

    let filter_json = try!(json_object.get("filter").ok_or(QueryParseError::FilteredNoFilter));
    let filter = match parse_filter(filter_json) {
        Ok(filter) => filter,
        Err(err) => return Err(QueryParseError::FilterParseError(err)),
    };

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

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

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

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

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

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

        assert_eq!(query, Ok(Query::MatchAll))
    }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


<commit_before>//! Defines complex queries over Ardite driver data structures.

use std::convert::From;
use linear_map::LinearMap;
use value::{Key, Pointer, Value};

/// Specifies a positive numeric range of data.
pub struct Range(Bound<u64>, Bound<u64>);

/// Taken from unstable [Rust][1]. When stable we should use it.
///
/// [1]: http://doc.rust-lang.org/std/collections/enum.Bound.html
pub enum Bound<T> {
  Included(T),
  Excluded(T),
  Unbounded
}

/// A condition which will resolve to a boolean value after comparing a certain
/// value with a set rule.
// TODO: Add standard conditions.
pub enum Condition {
  /// The condition always passes.
  True,
  /// The condition always fails.
  False,
  /// Inverts a condition.
  Not(Box<Condition>),
  /// Composes many conditions. They all must be true for the condition to be
  /// true.
  And(Vec<Condition>),
  /// Composes many conditions. Only one must be true for the condition to be
  /// true.
  Or(Vec<Condition>),
  /// If the compared value is exactly equal to this one, the condition passes.
  Equal(Value)
}

/// Specifies a complex driver query. The query is structured like a tree
/// except each node is unaware of its name (or if it even has a name). It
/// cannot be expected that a `Query` tree will map 1 to 1 with a `Value` tree.
#[derive(PartialEq, Debug)]
pub enum Query {
  /// Queries a single value.
  All,
  /// Queries some partial properties of an object.
  Keys(LinearMap<Key, Query>)
}

impl From<Pointer> for Query {
  fn from(pointer: Pointer) -> Self {
    // Reverse loop through the pointer to construct the query.
    pointer.iter().rev().fold(Query::All, |acc, key| {
      let mut map = LinearMap::new();
      map.insert(key.to_owned(), acc);
      Query::Keys(map)
    })
  }
}

#[cfg(test)]
mod tests {
  use query::Query;

  #[test]
  fn test_from_pointer() {
    assert_eq!(Query::from(point!["hello", "good", "world"]), Query::Keys(linear_map! {
      S!("hello") => Query::Keys(linear_map! {
        S!("good") => Query::Keys(linear_map! {
          S!("world") => Query::All
        })
      })
    }));
    assert_eq!(Query::from(point!["good"]), Query::Keys(linear_map! {
      S!("good") => Query::All
    }));
    assert_eq!(Query::from(point![]), Query::All);
  }
}
<commit_msg>refactor: remove bound generics<commit_after>//! Defines complex queries over Ardite driver data structures.

use std::convert::From;
use linear_map::LinearMap;
use value::{Key, Pointer, Value};

/// Specifies a positive numeric range of data.
pub struct Range(Bound, Bound);

/// Taken from unstable [Rust][1]. When stable we should use it.
///
/// [1]: http://doc.rust-lang.org/std/collections/enum.Bound.html
pub enum Bound {
  Included(u64),
  Excluded(u64),
  Unbounded
}

/// A condition which will resolve to a boolean value after comparing a certain
/// value with a set rule.
// TODO: Add standard conditions.
pub enum Condition {
  /// The condition always passes.
  True,
  /// The condition always fails.
  False,
  /// Inverts a condition.
  Not(Box<Condition>),
  /// Composes many conditions. They all must be true for the condition to be
  /// true.
  And(Vec<Condition>),
  /// Composes many conditions. Only one must be true for the condition to be
  /// true.
  Or(Vec<Condition>),
  /// If the compared value is exactly equal to this one, the condition passes.
  Equal(Value)
}

/// Specifies a complex driver query. The query is structured like a tree
/// except each node is unaware of its name (or if it even has a name). It
/// cannot be expected that a `Query` tree will map 1 to 1 with a `Value` tree.
#[derive(PartialEq, Debug)]
pub enum Query {
  /// Queries a single value.
  All,
  /// Queries some partial properties of an object.
  Keys(LinearMap<Key, Query>)
}

impl From<Pointer> for Query {
  fn from(pointer: Pointer) -> Self {
    // Reverse loop through the pointer to construct the query.
    pointer.iter().rev().fold(Query::All, |acc, key| {
      let mut map = LinearMap::new();
      map.insert(key.to_owned(), acc);
      Query::Keys(map)
    })
  }
}

#[cfg(test)]
mod tests {
  use query::Query;

  #[test]
  fn test_from_pointer() {
    assert_eq!(Query::from(point!["hello", "good", "world"]), Query::Keys(linear_map! {
      S!("hello") => Query::Keys(linear_map! {
        S!("good") => Query::Keys(linear_map! {
          S!("world") => Query::All
        })
      })
    }));
    assert_eq!(Query::from(point!["good"]), Query::Keys(linear_map! {
      S!("good") => Query::All
    }));
    assert_eq!(Query::from(point![]), Query::All);
  }
}
<|endoftext|>


<commit_before>use std::cmp;
use fnv::FnvHashMap;
use error::{Error, ErrorKind};
use result::Result;
use utf8::{DOLLAR, DOT};

const ROOT_QUERY_STR_OFFSET: usize = 2;

#[derive(Debug)]
pub struct Query<'a> {
    pub i: usize,
    pub ri: usize,
    pub target: bool,
    pub children: Option<FnvHashMap<&'a [u8], Query<'a>>>,
}

/// A pattern tree associated with the queries
pub struct QueryTree<'a> {
    pub root: FnvHashMap<&'a [u8], Query<'a>>,
    pub num_queries: usize,

    pub level: usize,
    pub qi: usize,
}

impl<'a> QueryTree<'a> {
    pub fn new<S: ?Sized + AsRef<[u8]>>(queries: &[&'a S]) -> Result<Self> {
        let mut root = FnvHashMap::default();
        let mut level = 0;
        let mut qi = 0;

        for (ri, s) in queries.into_iter().map(|s| (*s).as_ref()).enumerate() {
            if !is_valid_query_str(s) {
                return Err(Error::from(ErrorKind::InvalidQuery));
            }

            let l = set_queries(&mut root, &s[ROOT_QUERY_STR_OFFSET..], &mut qi, ri);
            level = cmp::max(level, l);
        }

        Ok(Self {
            root,
            num_queries: queries.len(),
            level,
            qi,
        })
    }
}

#[inline]
fn is_valid_query_str(query_str: &[u8]) -> bool {
    if query_str.len() < ROOT_QUERY_STR_OFFSET + 1 || query_str[0] != DOLLAR || query_str[1] != DOT {
        return false;
    }
    let mut s = ROOT_QUERY_STR_OFFSET - 1;
    for i in s + 1..query_str.len() {
        if query_str[i] != DOT {
            continue;
        }
        if i == s + 1 || i == query_str.len() - 1 {
            return false;
        }
        s = i;
    }
    true
}

#[inline]
fn set_queries<'a>(queries: &mut FnvHashMap<&'a [u8], Query<'a>>, s: &'a [u8], qi: &mut usize, ri: usize) -> usize {
    for j in 0..s.len() {
        if s[j] == DOT {
            let (t, u) = s.split_at(j);
            let query = queries.entry(t).or_insert_with(|| {
                let query = Query {
                    i: *qi,
                    ri: ri,
                    target: false,
                    children: None,
                };
                *qi += 1;
                query
            });
            let mut children = query.children.get_or_insert(FnvHashMap::default());
            return set_queries(&mut children, &u[1..], qi, ri) + 1;
        }
    }

    if !queries.contains_key(s) {
        queries.insert(
            s,
            Query {
                i: *qi,
                ri: ri,
                target: false,
                children: None,
            },
        );
        *qi += 1;
    }
    queries.get_mut(s).unwrap().target = true;
    1
}

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

    #[test]
    fn test_is_valid_query_str() {
        struct TestCase<'a> {
            query_str: &'a str,
            want: bool,
        }
        let test_cases = vec![
            TestCase {
                query_str: "",
                want: false,
            },
            TestCase {
                query_str: "$",
                want: false,
            },
            TestCase {
                query_str: "$.",
                want: false,
            },
            TestCase {
                query_str: "$..",
                want: false,
            },
            TestCase {
                query_str: "a.a",
                want: false,
            },
            TestCase {
                query_str: "$aa",
                want: false,
            },
            TestCase {
                query_str: "$.a",
                want: true,
            },
            TestCase {
                query_str: "$.aaaa",
                want: true,
            },
            TestCase {
                query_str: "$.aaaa.",
                want: false,
            },
            TestCase {
                query_str: "$.aaaa.b",
                want: true,
            },
            TestCase {
                query_str: "$.aaaa.bbbb",
                want: true,
            },
            TestCase {
                query_str: "$.aaaa.bbbb.",
                want: false,
            },
        ];
        for t in test_cases {
            let got = is_valid_query_str(t.query_str.as_bytes());
            assert_eq!(t.want, got);
        }
    }
}
<commit_msg>refactor(set_queries): simplify<commit_after>use std::cmp;
use fnv::FnvHashMap;
use error::{Error, ErrorKind};
use result::Result;
use utf8::{DOLLAR, DOT};

const ROOT_QUERY_STR_OFFSET: usize = 2;

#[derive(Debug)]
pub struct Query<'a> {
    pub i: usize,
    pub ri: usize,
    pub target: bool,
    pub children: Option<FnvHashMap<&'a [u8], Query<'a>>>,
}

/// A pattern tree associated with the queries
pub struct QueryTree<'a> {
    pub root: FnvHashMap<&'a [u8], Query<'a>>,
    pub num_queries: usize,

    pub level: usize,
    pub qi: usize,
}

impl<'a> QueryTree<'a> {
    pub fn new<S: ?Sized + AsRef<[u8]>>(queries: &[&'a S]) -> Result<Self> {
        let mut root = FnvHashMap::default();
        let mut level = 0;
        let mut qi = 0;

        for (ri, s) in queries.into_iter().map(|s| (*s).as_ref()).enumerate() {
            if !is_valid_query_str(s) {
                return Err(Error::from(ErrorKind::InvalidQuery));
            }

            let l = set_queries(&mut root, &s[ROOT_QUERY_STR_OFFSET..], &mut qi, ri);
            level = cmp::max(level, l);
        }

        Ok(Self {
            root,
            num_queries: queries.len(),
            level,
            qi,
        })
    }
}

#[inline]
fn is_valid_query_str(query_str: &[u8]) -> bool {
    if query_str.len() < ROOT_QUERY_STR_OFFSET + 1 || query_str[0] != DOLLAR || query_str[1] != DOT {
        return false;
    }
    let mut s = ROOT_QUERY_STR_OFFSET - 1;
    for i in s + 1..query_str.len() {
        if query_str[i] != DOT {
            continue;
        }
        if i == s + 1 || i == query_str.len() - 1 {
            return false;
        }
        s = i;
    }
    true
}

#[inline]
fn set_queries<'a>(queries: &mut FnvHashMap<&'a [u8], Query<'a>>, s: &'a [u8], qi: &mut usize, ri: usize) -> usize {
    for j in 0..s.len() {
        if s[j] == DOT {
            let (t, u) = s.split_at(j);
            let query = queries.entry(t).or_insert_with(|| {
                let query = Query {
                    i: *qi,
                    ri: ri,
                    target: false,
                    children: None,
                };
                *qi += 1;
                query
            });
            let mut children = query.children.get_or_insert(FnvHashMap::default());
            return set_queries(&mut children, &u[1..], qi, ri) + 1;
        }
    }

    let query = queries.entry(s).or_insert_with(|| {
        let query = Query {
            i: *qi,
            ri: ri,
            target: false,
            children: None,
        };
        *qi += 1;
        query
    });
    query.target = true;
    1
}

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

    #[test]
    fn test_is_valid_query_str() {
        struct TestCase<'a> {
            query_str: &'a str,
            want: bool,
        }
        let test_cases = vec![
            TestCase {
                query_str: "",
                want: false,
            },
            TestCase {
                query_str: "$",
                want: false,
            },
            TestCase {
                query_str: "$.",
                want: false,
            },
            TestCase {
                query_str: "$..",
                want: false,
            },
            TestCase {
                query_str: "a.a",
                want: false,
            },
            TestCase {
                query_str: "$aa",
                want: false,
            },
            TestCase {
                query_str: "$.a",
                want: true,
            },
            TestCase {
                query_str: "$.aaaa",
                want: true,
            },
            TestCase {
                query_str: "$.aaaa.",
                want: false,
            },
            TestCase {
                query_str: "$.aaaa.b",
                want: true,
            },
            TestCase {
                query_str: "$.aaaa.bbbb",
                want: true,
            },
            TestCase {
                query_str: "$.aaaa.bbbb.",
                want: false,
            },
        ];
        for t in test_cases {
            let got = is_valid_query_str(t.query_str.as_bytes());
            assert_eq!(t.want, got);
        }
    }
}
<|endoftext|>