<commit_before>//
// Copyright (c) Pirmin Kalberer. All rights reserved.
// Licensed under the MIT License. See LICENSE file in the project root for full license information.
//

use core::config::ApplicationCfg;
use datasource::postgis::PostgisInput;
use core::grid::Grid;
use service::mvt::{MvtService, Tileset};
use core::{Config, read_config, parse_config};
use core::config::DEFAULT_CONFIG;
use serde_json;
use cache::{Tilecache, Nocache, Filecache};

use nickel::{Nickel, Options, HttpRouter, MediaType, Request, Response, MiddlewareResult,
             StaticFilesHandler};
use hyper::header::{CacheControl, CacheDirective, AccessControlAllowOrigin,
                    AccessControlAllowMethods, ContentEncoding, Encoding};
use hyper::method::Method;
use hyper::header;
use std::collections::HashMap;
use std::str::FromStr;
use clap::ArgMatches;
use std::str;
use std::process;
use open;

fn log_request<'mw>(req: &mut Request<MvtService>,
                    res: Response<'mw, MvtService>)
                    -> MiddlewareResult<'mw, MvtService> {
    info!("{} {}", req.origin.method, req.origin.uri);
    res.next_middleware()
}

header! { (ContentType, "Content-Type") => [String] }

#[derive(RustcEncodable)]
struct TilesetInfo {
    name: String,
    layerinfos: String,
    hasviewer: bool,
}

impl TilesetInfo {
    fn from_tileset(set: &Tileset) -> TilesetInfo {
        let mut hasviewer = true;
        let layerinfos: Vec<String> = set.layers
            .iter()
            .map(|l| {
                let geom_type = l.geometry_type.clone().unwrap_or("UNKNOWN".to_string());
                hasviewer = hasviewer &&
                            ["POINT",
                             "LINESTRING",
                             "POLYGON",
                             "MULTPOINT",
                             "MULTILINESTRING",
                             "MULTIPOLYGON"]
                                    .contains(&(&geom_type as &str));
                format!("{} [{}]", &l.name, &geom_type)
            })
            .collect();
        TilesetInfo {
            name: set.name.clone(),
            layerinfos: format!("{}", layerinfos.join(", ")),
            hasviewer: hasviewer,
        }
    }
}

struct StaticFiles {
    files: HashMap<&'static str, (&'static [u8], MediaType)>,
}

impl StaticFiles {
    fn init() -> StaticFiles {
        let mut static_files = StaticFiles { files: HashMap::new() };
        static_files.add("favicon.ico",
                         include_bytes!("static/favicon.ico"),
                         MediaType::Ico);
        static_files.add("index.html",
                         include_bytes!("static/index.html"),
                         MediaType::Html);
        static_files.add("viewer.js",
                         include_bytes!("static/viewer.js"),
                         MediaType::Js);
        static_files.add("viewer.css",
                         include_bytes!("static/viewer.css"),
                         MediaType::Css);
        static_files.add("maputnik.html",
                         include_bytes!("static/maputnik.html"),
                         MediaType::Html);
        static_files.add("maputnik.js",
                         include_bytes!("static/maputnik.js"),
                         MediaType::Js);
        static_files.add("maputnik-vendor.js",
                         include_bytes!("static/maputnik-vendor.js"),
                         MediaType::Js);
        static_files.add("img/maputnik.png",
                         include_bytes!("static/img/maputnik.png"),
                         MediaType::Png);
        static_files.add("fonts/Roboto-Regular.ttf",
                         include_bytes!("static/fonts/Roboto-Regular.ttf"),
                         MediaType::Ttf);
        static_files.add("fonts/Roboto-Medium.ttf",
                         include_bytes!("static/fonts/Roboto-Medium.ttf"),
                         MediaType::Ttf);
        static_files
    }
    fn add(&mut self, name: &'static str, data: &'static [u8], media_type: MediaType) {
        self.files.insert(name, (data, media_type));
    }
    fn content(&self, base: Option<&str>, name: String) -> Option<&(&[u8], MediaType)> {
        let mut key = if name == "." {
            "index.html".to_string()
        } else {
            name
        };
        if let Some(path) = base {
            key = format!("{}/{}", path, key);
        }
        self.files.get(&key as &str)
    }
}

pub fn service_from_args(args: &ArgMatches) -> (MvtService, ApplicationCfg) {
    if let Some(cfgpath) = args.value_of("config") {
        info!("Reading configuration from '{}'", cfgpath);
        let config =
            read_config(cfgpath).unwrap_or_else(|err| {
                                                    println!("Error reading configuration - {} ",
                                                             err);
                                                    process::exit(1)
                                                });
        let mut svc =
            MvtService::from_config(&config).unwrap_or_else(|err| {
                                                                println!("Error reading configuration - {} ", err);
                                                                process::exit(1)
                                                            });
        svc.connect();
        (svc, config)
    } else {
        let config = parse_config(DEFAULT_CONFIG.to_string(), "").unwrap();
        let cache = match args.value_of("cache") {
            None => Tilecache::Nocache(Nocache),
            Some(dir) => {
                Tilecache::Filecache(Filecache {
                                         basepath: dir.to_string(),
                                         baseurl: None,
                                     })
            }
        };
        let simplify = bool::from_str(args.value_of("simplify").unwrap_or("true")).unwrap_or(false);
        let clip = bool::from_str(args.value_of("clip").unwrap_or("true")).unwrap_or(false);
        if let Some(dbconn) = args.value_of("dbconn") {
            let pg = PostgisInput::new(dbconn).connected();
            let grid = Grid::web_mercator();
            let detect_geometry_types = true; //TODO: add option (maybe slow for many geometries)
            let mut layers = pg.detect_layers(detect_geometry_types);
            let mut tilesets = Vec::new();
            while let Some(mut l) = layers.pop() {
                let extent = pg.layer_extent(&l);
                l.simplify = Some(simplify);
                if simplify {
                    // Limit features by default unless simplify is set to false
                    l.query_limit = Some(1000);
                }
                l.buffer_size = match l.geometry_type {
                    Some(ref geom) => {
                        if clip {
                            if geom.contains("POLYGON") {
                                Some(1)
                            } else {
                                Some(0)
                            }
                        } else {
                            None
                        }
                    }
                    None => None,
                };
                let tileset = Tileset {
                    name: l.name.clone(),
                    extent: extent,
                    layers: vec![l],
                };
                tilesets.push(tileset);
            }
            let svc = MvtService {
                input: pg,
                grid: grid,
                tilesets: tilesets,
                cache: cache,
            };
            (svc, config)
        } else {
            println!("Either 'config' or 'dbconn' is required");
            process::exit(1)
        }
    }
}

pub fn webserver(args: &ArgMatches) {
    let (mut service, config) = service_from_args(args);

    let mvt_viewer = config.service.mvt.viewer;
    let bind: &str = &config.webserver.bind.unwrap_or("127.0.0.1".to_string());
    let port = config.webserver.port.unwrap_or(6767);
    let threads = config.webserver.threads.unwrap_or(4) as usize;
    let cache_max_age = config.webserver.cache_control_max_age.unwrap_or(300);

    service.prepare_feature_queries();
    service.init_cache();

    let mut tileset_infos: Vec<TilesetInfo> = service
        .tilesets
        .iter()
        .map(|set| TilesetInfo::from_tileset(&set))
        .collect();
    tileset_infos.sort_by_key(|ti| ti.name.clone());

    let mut server = Nickel::with_data(service);
    server.options = Options::default().thread_count(Some(threads));
    // Avoid thread exhaustion caused by hypers keep_alive handling (https://github.com/hyperium/hyper/issues/368)
    server.keep_alive_timeout(None);
    server.utilize(log_request);

    server.get("/**(.style)?.json",
               middleware! { |_req, mut res|
        res.set(MediaType::Json);
        res.set(AccessControlAllowMethods(vec![Method::Get]));
        res.set(AccessControlAllowOrigin::Any);
    });

    server.get("/index.json",
               middleware! { |_req, res|
        let service: &MvtService = res.server_data();
        let json = service.get_mvt_metadata().unwrap();
        serde_json::to_vec(&json).unwrap()
    });

    // Font list for Maputnik
    server.get("/fontstacks.json",
               middleware! { |_req, _res|
        let json = json!(["Roboto Medium","Roboto Regular"]);
        serde_json::to_vec(&json).unwrap()
    });

    server.get("/:tileset.json",
               middleware! { |req, res|
        let service: &MvtService = res.server_data();
        let tileset = req.param("tileset").unwrap();
        let host = req.origin.headers.get::<header::Host>().unwrap();
        let baseurl = format!("http://{}:{}", host.hostname, host.port.unwrap_or(80));
        let json = service.get_tilejson(&baseurl, &tileset).unwrap();
        serde_json::to_vec(&json).unwrap()
    });

    server.get("/:tileset.style.json",
               middleware! { |req, res|
        let service: &MvtService = res.server_data();
        let tileset = req.param("tileset").unwrap();
        let host = req.origin.headers.get::<header::Host>().unwrap();
        let baseurl = format!("http://{}:{}", host.hostname, host.port.unwrap_or(80));
        let json = service.get_stylejson(&baseurl, &tileset).unwrap();
        serde_json::to_vec(&json).unwrap()
    });

    server.get("/:tileset/metadata.json",
               middleware! { |req, res|
        let service: &MvtService = res.server_data();
        let tileset = req.param("tileset").unwrap();
        let json = service.get_mbtiles_metadata(&tileset).unwrap();
        serde_json::to_vec(&json).unwrap()
    });

    server.get("/:tileset/:z/:x/:y.pbf",
               middleware! { |req, mut res|
        let service: &MvtService = res.server_data();

        let tileset = req.param("tileset").unwrap();
        let z = req.param("z").unwrap().parse::<u8>().unwrap();
        let x = req.param("x").unwrap().parse::<u32>().unwrap();
        let y = req.param("y").unwrap().parse::<u32>().unwrap();

        let accept_encoding = req.origin.headers.get::<header::AcceptEncoding>();
        let gzip = accept_encoding.is_some() && accept_encoding.unwrap().iter().any(
                   |ref qit| qit.item == Encoding::Gzip );
        let tile = service.tile_cached(tileset, x, y, z, gzip);
        if gzip {
            res.set_header_fallback(|| ContentEncoding(vec![Encoding::Gzip]));
        }
        res.set_header_fallback(|| ContentType("application/x-protobuf".to_owned()));
        res.set_header_fallback(|| CacheControl(vec![CacheDirective::MaxAge(cache_max_age)]));
        //res.set_header_fallback(|| ContentLength(tile.len() as u64));
        res.set(AccessControlAllowMethods(vec![Method::Get]));
        res.set(AccessControlAllowOrigin::Any);

        tile
    });

    if mvt_viewer {
        let static_files = StaticFiles::init();
        server.get("/(:base/)?:static",
                   middleware! { |req, mut res|
            let mut name = req.param("static").unwrap().to_string();
            if let Some(format) = req.param("format") {
                name = format!("{}.{}", name, format);
            }
            if let Some(content) = static_files.content(req.param("base"), name) {
                res.set(content.1);
                return res.send(content.0)
            }
        });
    }

    server.get("/**", StaticFilesHandler::new("public/"));

    let _listening = server
        .listen((bind, port))
        .expect("Failed to launch server");

    let openbrowser = bool::from_str(args.value_of("openbrowser").unwrap_or("true"))
        .unwrap_or(false);
    if openbrowser && mvt_viewer {
        let _res = open::that(format!("http://{}:{}", bind, port));
    }
}

pub fn gen_config(args: &ArgMatches) -> String {
    let toml = r#"
[webserver]
# Bind address. Use 0.0.0.0 to listen on all adresses.
bind = "127.0.0.1"
port = 6767
threads = 4
#cache_control_max_age = 43200
"#;
    let mut config;
    if let Some(_dbconn) = args.value_of("dbconn") {
        let (service, _) = service_from_args(args);
        config = service.gen_runtime_config();
    } else {
        config = MvtService::gen_config();
    }
    config.push_str(toml);
    config
}

#[test]
fn test_gen_config() {
    use core::parse_config;

    let args = ArgMatches::new();
    let toml = gen_config(&args);
    println!("{}", toml);
    assert_eq!(Some("# t-rex configuration"), toml.lines().next());

    let config = parse_config(toml, "").unwrap();
    let service = MvtService::from_config(&config).unwrap();
    assert_eq!(service.input.connection_url,
               "postgresql://user:pass@host/database");
}

#[test]
#[ignore]
fn test_runtime_config() {
    use std::env;
    use clap::App;
    use core::parse_config;

    if env::var("DBCONN").is_err() {
        panic!("DBCONN undefined");
    }
    let args = App::new("test")
        .args_from_usage("--dbconn=[SPEC] 'PostGIS connection postgresql://USER@HOST/DBNAME'")
        .get_matches_from(vec!["", "--dbconn", &env::var("DBCONN").unwrap()]);
    let toml = gen_config(&args);
    println!("{}", toml);
    assert_eq!(Some("# t-rex configuration"), toml.lines().next());

    let config = parse_config(toml, "").unwrap();
    let service = MvtService::from_config(&config).unwrap();
    assert_eq!(service.input.connection_url, env::var("DBCONN").unwrap());
}
<commit_msg>Show dino<commit_after>//
// Copyright (c) Pirmin Kalberer. All rights reserved.
// Licensed under the MIT License. See LICENSE file in the project root for full license information.
//

use core::config::ApplicationCfg;
use datasource::postgis::PostgisInput;
use core::grid::Grid;
use service::mvt::{MvtService, Tileset};
use core::{Config, read_config, parse_config};
use core::config::DEFAULT_CONFIG;
use serde_json;
use cache::{Tilecache, Nocache, Filecache};

use nickel::{Nickel, Options, HttpRouter, MediaType, Request, Response, MiddlewareResult,
             StaticFilesHandler};
use hyper::header::{CacheControl, CacheDirective, AccessControlAllowOrigin,
                    AccessControlAllowMethods, ContentEncoding, Encoding};
use hyper::method::Method;
use hyper::header;
use std::collections::HashMap;
use std::str::FromStr;
use clap::ArgMatches;
use std::str;
use std::process;
use open;

fn log_request<'mw>(req: &mut Request<MvtService>,
                    res: Response<'mw, MvtService>)
                    -> MiddlewareResult<'mw, MvtService> {
    info!("{} {}", req.origin.method, req.origin.uri);
    res.next_middleware()
}

header! { (ContentType, "Content-Type") => [String] }

#[derive(RustcEncodable)]
struct TilesetInfo {
    name: String,
    layerinfos: String,
    hasviewer: bool,
}

impl TilesetInfo {
    fn from_tileset(set: &Tileset) -> TilesetInfo {
        let mut hasviewer = true;
        let layerinfos: Vec<String> = set.layers
            .iter()
            .map(|l| {
                let geom_type = l.geometry_type.clone().unwrap_or("UNKNOWN".to_string());
                hasviewer = hasviewer &&
                            ["POINT",
                             "LINESTRING",
                             "POLYGON",
                             "MULTPOINT",
                             "MULTILINESTRING",
                             "MULTIPOLYGON"]
                                    .contains(&(&geom_type as &str));
                format!("{} [{}]", &l.name, &geom_type)
            })
            .collect();
        TilesetInfo {
            name: set.name.clone(),
            layerinfos: format!("{}", layerinfos.join(", ")),
            hasviewer: hasviewer,
        }
    }
}

struct StaticFiles {
    files: HashMap<&'static str, (&'static [u8], MediaType)>,
}

impl StaticFiles {
    fn init() -> StaticFiles {
        let mut static_files = StaticFiles { files: HashMap::new() };
        static_files.add("favicon.ico",
                         include_bytes!("static/favicon.ico"),
                         MediaType::Ico);
        static_files.add("index.html",
                         include_bytes!("static/index.html"),
                         MediaType::Html);
        static_files.add("viewer.js",
                         include_bytes!("static/viewer.js"),
                         MediaType::Js);
        static_files.add("viewer.css",
                         include_bytes!("static/viewer.css"),
                         MediaType::Css);
        static_files.add("maputnik.html",
                         include_bytes!("static/maputnik.html"),
                         MediaType::Html);
        static_files.add("maputnik.js",
                         include_bytes!("static/maputnik.js"),
                         MediaType::Js);
        static_files.add("maputnik-vendor.js",
                         include_bytes!("static/maputnik-vendor.js"),
                         MediaType::Js);
        static_files.add("img/maputnik.png",
                         include_bytes!("static/img/maputnik.png"),
                         MediaType::Png);
        static_files.add("fonts/Roboto-Regular.ttf",
                         include_bytes!("static/fonts/Roboto-Regular.ttf"),
                         MediaType::Ttf);
        static_files.add("fonts/Roboto-Medium.ttf",
                         include_bytes!("static/fonts/Roboto-Medium.ttf"),
                         MediaType::Ttf);
        static_files
    }
    fn add(&mut self, name: &'static str, data: &'static [u8], media_type: MediaType) {
        self.files.insert(name, (data, media_type));
    }
    fn content(&self, base: Option<&str>, name: String) -> Option<&(&[u8], MediaType)> {
        let mut key = if name == "." {
            "index.html".to_string()
        } else {
            name
        };
        if let Some(path) = base {
            key = format!("{}/{}", path, key);
        }
        self.files.get(&key as &str)
    }
}

static DINO: &'static str = "             xxxxxxxxx
        xxxxxxxxxxxxxxxxxxxxxxxx
      xxxxxxxxxxxxxxxxxxxxxxxxxxxx
     xxxxxxxxxxxxxxxxxxxxxxxxx xxxx
     xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
    xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
    xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
   xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
 xxxxxxxxxxxxxxxxxxxxxxxxxxxxxx xxxxxxxxxxxxxx
xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx  xxxxxxxxxxxxxx
xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx   xxxxxxxxxxxxx
xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx   xxxxxxxxxx
xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx     xxxxxx
xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx      x
xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
xxxxxxxxxxxxxxxxxxxxxxxxxx    xxxxxxxxxxx
xxxxxxxxxxxxxx                   xxxxxx
xxxxxxxxxxxx
xxxxxxxxxxx
xxxxxxxxxx
xxxxxxxxx
xxxxxxx
xxxxxx
xxxxxxx";

pub fn service_from_args(args: &ArgMatches) -> (MvtService, ApplicationCfg) {
    if let Some(cfgpath) = args.value_of("config") {
        info!("Reading configuration from '{}'", cfgpath);
        let config =
            read_config(cfgpath).unwrap_or_else(|err| {
                                                    println!("Error reading configuration - {} ",
                                                             err);
                                                    process::exit(1)
                                                });
        let mut svc =
            MvtService::from_config(&config).unwrap_or_else(|err| {
                                                                println!("Error reading configuration - {} ", err);
                                                                process::exit(1)
                                                            });
        svc.connect();
        (svc, config)
    } else {
        let config = parse_config(DEFAULT_CONFIG.to_string(), "").unwrap();
        let cache = match args.value_of("cache") {
            None => Tilecache::Nocache(Nocache),
            Some(dir) => {
                Tilecache::Filecache(Filecache {
                                         basepath: dir.to_string(),
                                         baseurl: None,
                                     })
            }
        };
        let simplify = bool::from_str(args.value_of("simplify").unwrap_or("true")).unwrap_or(false);
        let clip = bool::from_str(args.value_of("clip").unwrap_or("true")).unwrap_or(false);
        if let Some(dbconn) = args.value_of("dbconn") {
            let pg = PostgisInput::new(dbconn).connected();
            let grid = Grid::web_mercator();
            let detect_geometry_types = true; //TODO: add option (maybe slow for many geometries)
            let mut layers = pg.detect_layers(detect_geometry_types);
            let mut tilesets = Vec::new();
            while let Some(mut l) = layers.pop() {
                let extent = pg.layer_extent(&l);
                l.simplify = Some(simplify);
                if simplify {
                    // Limit features by default unless simplify is set to false
                    l.query_limit = Some(1000);
                }
                l.buffer_size = match l.geometry_type {
                    Some(ref geom) => {
                        if clip {
                            if geom.contains("POLYGON") {
                                Some(1)
                            } else {
                                Some(0)
                            }
                        } else {
                            None
                        }
                    }
                    None => None,
                };
                let tileset = Tileset {
                    name: l.name.clone(),
                    extent: extent,
                    layers: vec![l],
                };
                tilesets.push(tileset);
            }
            let svc = MvtService {
                input: pg,
                grid: grid,
                tilesets: tilesets,
                cache: cache,
            };
            (svc, config)
        } else {
            println!("Either 'config' or 'dbconn' is required");
            process::exit(1)
        }
    }
}

pub fn webserver(args: &ArgMatches) {
    let (mut service, config) = service_from_args(args);

    let mvt_viewer = config.service.mvt.viewer;
    let bind: &str = &config.webserver.bind.unwrap_or("127.0.0.1".to_string());
    let port = config.webserver.port.unwrap_or(6767);
    let threads = config.webserver.threads.unwrap_or(4) as usize;
    let cache_max_age = config.webserver.cache_control_max_age.unwrap_or(300);

    service.prepare_feature_queries();
    service.init_cache();

    let mut tileset_infos: Vec<TilesetInfo> = service
        .tilesets
        .iter()
        .map(|set| TilesetInfo::from_tileset(&set))
        .collect();
    tileset_infos.sort_by_key(|ti| ti.name.clone());

    let mut server = Nickel::with_data(service);
    server.options = Options::default().thread_count(Some(threads));
    // Avoid thread exhaustion caused by hypers keep_alive handling (https://github.com/hyperium/hyper/issues/368)
    server.keep_alive_timeout(None);
    server.utilize(log_request);

    server.get("/**(.style)?.json",
               middleware! { |_req, mut res|
        res.set(MediaType::Json);
        res.set(AccessControlAllowMethods(vec![Method::Get]));
        res.set(AccessControlAllowOrigin::Any);
    });

    server.get("/index.json",
               middleware! { |_req, res|
        let service: &MvtService = res.server_data();
        let json = service.get_mvt_metadata().unwrap();
        serde_json::to_vec(&json).unwrap()
    });

    // Font list for Maputnik
    server.get("/fontstacks.json",
               middleware! { |_req, _res|
        let json = json!(["Roboto Medium","Roboto Regular"]);
        serde_json::to_vec(&json).unwrap()
    });

    server.get("/:tileset.json",
               middleware! { |req, res|
        let service: &MvtService = res.server_data();
        let tileset = req.param("tileset").unwrap();
        let host = req.origin.headers.get::<header::Host>().unwrap();
        let baseurl = format!("http://{}:{}", host.hostname, host.port.unwrap_or(80));
        let json = service.get_tilejson(&baseurl, &tileset).unwrap();
        serde_json::to_vec(&json).unwrap()
    });

    server.get("/:tileset.style.json",
               middleware! { |req, res|
        let service: &MvtService = res.server_data();
        let tileset = req.param("tileset").unwrap();
        let host = req.origin.headers.get::<header::Host>().unwrap();
        let baseurl = format!("http://{}:{}", host.hostname, host.port.unwrap_or(80));
        let json = service.get_stylejson(&baseurl, &tileset).unwrap();
        serde_json::to_vec(&json).unwrap()
    });

    server.get("/:tileset/metadata.json",
               middleware! { |req, res|
        let service: &MvtService = res.server_data();
        let tileset = req.param("tileset").unwrap();
        let json = service.get_mbtiles_metadata(&tileset).unwrap();
        serde_json::to_vec(&json).unwrap()
    });

    server.get("/:tileset/:z/:x/:y.pbf",
               middleware! { |req, mut res|
        let service: &MvtService = res.server_data();

        let tileset = req.param("tileset").unwrap();
        let z = req.param("z").unwrap().parse::<u8>().unwrap();
        let x = req.param("x").unwrap().parse::<u32>().unwrap();
        let y = req.param("y").unwrap().parse::<u32>().unwrap();

        let accept_encoding = req.origin.headers.get::<header::AcceptEncoding>();
        let gzip = accept_encoding.is_some() && accept_encoding.unwrap().iter().any(
                   |ref qit| qit.item == Encoding::Gzip );
        let tile = service.tile_cached(tileset, x, y, z, gzip);
        if gzip {
            res.set_header_fallback(|| ContentEncoding(vec![Encoding::Gzip]));
        }
        res.set_header_fallback(|| ContentType("application/x-protobuf".to_owned()));
        res.set_header_fallback(|| CacheControl(vec![CacheDirective::MaxAge(cache_max_age)]));
        //res.set_header_fallback(|| ContentLength(tile.len() as u64));
        res.set(AccessControlAllowMethods(vec![Method::Get]));
        res.set(AccessControlAllowOrigin::Any);

        tile
    });

    if mvt_viewer {
        let static_files = StaticFiles::init();
        server.get("/(:base/)?:static",
                   middleware! { |req, mut res|
            let mut name = req.param("static").unwrap().to_string();
            if let Some(format) = req.param("format") {
                name = format!("{}.{}", name, format);
            }
            if let Some(content) = static_files.content(req.param("base"), name) {
                res.set(content.1);
                return res.send(content.0)
            }
        });
    }

    server.get("/**", StaticFilesHandler::new("public/"));

    println!("{}", DINO);

    let _listening = server
        .listen((bind, port))
        .expect("Failed to launch server");

    let openbrowser = bool::from_str(args.value_of("openbrowser").unwrap_or("true"))
        .unwrap_or(false);
    if openbrowser && mvt_viewer {
        let _res = open::that(format!("http://{}:{}", bind, port));
    }
}

pub fn gen_config(args: &ArgMatches) -> String {
    let toml = r#"
[webserver]
# Bind address. Use 0.0.0.0 to listen on all adresses.
bind = "127.0.0.1"
port = 6767
threads = 4
#cache_control_max_age = 43200
"#;
    let mut config;
    if let Some(_dbconn) = args.value_of("dbconn") {
        let (service, _) = service_from_args(args);
        config = service.gen_runtime_config();
    } else {
        config = MvtService::gen_config();
    }
    config.push_str(toml);
    config
}

#[test]
fn test_gen_config() {
    use core::parse_config;

    let args = ArgMatches::new();
    let toml = gen_config(&args);
    println!("{}", toml);
    assert_eq!(Some("# t-rex configuration"), toml.lines().next());

    let config = parse_config(toml, "").unwrap();
    let service = MvtService::from_config(&config).unwrap();
    assert_eq!(service.input.connection_url,
               "postgresql://user:pass@host/database");
}

#[test]
#[ignore]
fn test_runtime_config() {
    use std::env;
    use clap::App;
    use core::parse_config;

    if env::var("DBCONN").is_err() {
        panic!("DBCONN undefined");
    }
    let args = App::new("test")
        .args_from_usage("--dbconn=[SPEC] 'PostGIS connection postgresql://USER@HOST/DBNAME'")
        .get_matches_from(vec!["", "--dbconn", &env::var("DBCONN").unwrap()]);
    let toml = gen_config(&args);
    println!("{}", toml);
    assert_eq!(Some("# t-rex configuration"), toml.lines().next());

    let config = parse_config(toml, "").unwrap();
    let service = MvtService::from_config(&config).unwrap();
    assert_eq!(service.input.connection_url, env::var("DBCONN").unwrap());
}
<|endoftext|>


<commit_before>//! Middle layer providing a somewhat safer (but still quite unsafe)
//! API.
//!
//! The main idea of the middle layer is to wrap types
//! [`ffi_cif`](../raw/struct.ffi_cif.html) and
//! [`ffi_closure`](../raw/struct.ffi_closure.html) as
//! [`Cif`](struct.Cif.html) and [`Closure`](struct.Closure.html),
//! respectively, so that their resources are managed properly. However,
//! calling a function via a CIF or closure is still unsafe because
//! argument types aren’t checked. See the [`high`](../high/index.html)
//! layer for closures with type-checked arguments.
//!
use std::os::raw::c_void; use std::marker::PhantomData;

use low;
pub use low::{Callback, CallbackMut, CodePtr,
              ffi_abi as FfiAbi, FFI_DEFAULT_ABI};

pub mod types;
pub use self::types::Type;

/// Contains an untyped pointer to a function argument.
///
/// When calling a function via a [CIF](struct.Cif.html), each argument
/// must be passed as a C `void*`. Wrapping the argument in the `Arg`
/// struct accomplishes the necessary coercion.
#[derive(Debug)]
#[repr(C)]
pub struct Arg(*mut c_void);

impl Arg {
    /// Coerces an argument reference into the `Arg` type.
    ///
    /// This is used to wrap each argument pointer before passing them
    /// to [`Cif::call`](struct.Cif.html#method.call).
    pub fn new<T>(r: &T) -> Self {
        Arg(r as *const T as *mut c_void)
    }
}

/// Coerces an argument reference into the [`Arg`](struct.Arg.html)
/// type.
///
/// This is used to wrap each argument pointer before passing them
/// to [`Cif::call`](struct.Cif.html#method.call).
/// (This is the same as [`Arg::new`](struct.Arg.html#method.new)).
pub fn arg<T>(r: &T) -> Arg {
    Arg::new(r)
}

/// Describes the calling convention and types for calling a function.
///
/// This is the `middle` layer’s wrapping of the `low` and `raw` layers’
/// [`ffi_cif`](../raw/struct.ffi_cif.html). An initialized CIF contains
/// references to an array of argument types and a result type, each of
/// which may be allocated on the heap. `Cif` manages the memory of
/// those referenced objects.
///
/// Construct with [`Cif::new`](#method.new) or
/// [`Cif::from_type_array`](#method.from_type_array).
///
/// # Example
///
/// ```
/// extern "C" fn add(x: f64, y: &f64) -> f64 {
///     return x + y;
/// }
///
/// use libffi::middle::*;
///
/// let args = vec![Type::f64(), Type::pointer()];
/// let cif = Cif::new(args.into_iter(), Type::f64());
///
/// let n = unsafe { cif.call(CodePtr(add as _), &[arg(&5), arg(&&6)]) };
/// assert_eq!(11, n);
/// ```
#[derive(Clone, Debug)]
pub struct Cif {
    cif:    low::ffi_cif,
    args:   types::TypeArray,
    result: Type,
}

impl Cif {
    /// Creates a new CIF for the given argument and result types.
    ///
    /// Takes ownership of the argument and result
    /// [`Type`](types/struct.Type.html)s, because the resulting
    /// `Cif` retains references to them.
    /// Defaults to the platform’s default calling convention; this
    /// can be adjusted using [`set_abi`](#method.set_abi).
    pub fn new<I>(args: I, result: Type) -> Self
        where I: ExactSizeIterator<Item=Type>
    {
        Self::from_type_array(types::TypeArray::new(args), result)
    }

    /// Calls a function with the given arguments.
    ///
    /// In particular, this method invokes function `fun` passing it
    /// arguments `args`, and returns the result.
    ///
    /// # Safety
    ///
    /// There is no checking that the calling convention and types
    /// in the `Cif` match the actual calling convention and type of
    /// `fun`.
    pub unsafe fn call<R>(&self, fun: CodePtr, args: &[Arg]) -> R {
        use std::mem;

        assert!(self.cif.nargs as usize == args.len(),
                "Cif::call: passed wrong number of arguments");

        low::call::<R>(&self.cif as *const _ as *mut _,
                       fun,
                       mem::transmute::<*const Arg,
                                        *mut *mut c_void>(args.as_ptr()))
    }

    /// Sets the CIF to use the given calling convention.
    pub fn set_abi(&mut self, abi: FfiAbi) {
        self.cif.abi = abi;
    }

    /// Creates a new CIF for the given argument and result types.
    ///
    /// This is just like [`Cif::new`](#method.new), except it takes a
    /// `TypeArray` instead of an `ExactSizeIterator`.
    pub fn from_type_array(args: types::TypeArray, result: Type) -> Self {
        let mut cif: low::ffi_cif = Default::default();

        unsafe {
            low::prep_cif(&mut cif,
                          low::FFI_DEFAULT_ABI,
                          args.len(),
                          result.as_raw_ptr(),
                          args.as_raw_ptr())
        }.expect("low::prep_cif");

        // Note that cif retains references to args and result,
        // which is why we hold onto them here.
        Cif {
            cif:    cif,
            args:   args,
            result: result,
        }
    }

    /// Gets a raw pointer to the underlying
    /// [`ffi_cif`](../low/struct.ffi_cif.html).
    ///
    /// This can be used for passing a `middle::Cif` to functions from the
    /// [`low`](../low/index.html) and [`raw`](../raw/index.html) modules.
    pub fn as_raw_ptr(&self) -> *mut low::ffi_cif {
        &self.cif as *const _ as *mut _
    }
}

/// Represents a closure callable from C.
///
/// A libffi closure captures a `void*` (“userdata”) and passes it to a
/// callback when the code pointer (obtained via
/// [`code_ptr`](#method.code_ptr)) is invoked.
///
/// Construct with [`Closure::new`](#method.new) and
/// [`Closure::new_mut`](#method.new_mut).
///
/// # Example
///
/// In this example we turn a Rust lambda into a C function. Note that
/// the Rust closure value `lambda` is actually passed to the callback
/// function as userdata, which then invokes it.
///
/// ```
/// use std::mem;
/// use std::os::raw::c_void;
///
/// use libffi::middle::*;
/// use libffi::low;
///
/// unsafe extern "C" fn lambda_callback<F: Fn(u64, u64) -> u64>(
///     _cif: &low::ffi_cif,
///     result: &mut u64,
///     args: *const *const c_void,
///     userdata: &F)
/// {
///     let args: *const &u64 = mem::transmute(args);
///     let arg1 = **args.offset(0);
///     let arg2 = **args.offset(1);
///
///     *result = userdata(arg1, arg2);
/// }
///
/// fn rust_lambda() {
///     let cif = Cif::new(vec![Type::u64(), Type::u64()].into_iter(),
///                        Type::u64());
///     let lambda = |x: u64, y: u64| x + y;
///     let closure = Closure::new(cif, lambda_callback, &lambda);
///
///     unsafe {
///         let fun: &unsafe extern "C" fn(u64, u64) -> u64
///             = mem::transmute(closure.code_ptr());
///
///         assert_eq!(11, fun(5, 6));
///         assert_eq!(12, fun(5, 7));
///     }
/// }
/// ```
#[derive(Debug)]
pub struct Closure<'a> {
    _cif:    Box<Cif>,
    alloc:   *mut ::low::ffi_closure,
    code:    CodePtr,
    _marker: PhantomData<&'a ()>,
}

impl<'a> Drop for Closure<'a> {
    fn drop(&mut self) {
        unsafe {
            low::closure_free(self.alloc);
        }
    }
}

impl<'a> Closure<'a> {
    /// Creates a new closure with immutable userdata.
    ///
    /// # Arguments
    ///
    /// - `cif` — describes the calling convention and argument and
    ///   result types
    /// - `callback` — the function to call when the closure is invoked
    /// - `userdata` — the pointer to pass to `callback` along with the
    ///   arguments when the closure is called
    ///
    /// # Result
    ///
    /// The new closure.
    pub fn new<U, R>(cif:      Cif,
                     callback: Callback<U, R>,
                     userdata: &'a U) -> Self
    {
        let cif = Box::new(cif);
        let (alloc, code) = low::closure_alloc();

        unsafe {
            low::prep_closure(alloc,
                              cif.as_raw_ptr(),
                              callback,
                              userdata as *const U,
                              code).unwrap();
        }

        Closure {
            _cif:    cif,
            alloc:   alloc,
            code:    code,
            _marker: PhantomData,
        }
    }

    /// Creates a new closure with mutable userdata.
    ///
    /// # Arguments
    ///
    /// - `cif` — describes the calling convention and argument and
    ///   result types
    /// - `callback` — the function to call when the closure is invoked
    /// - `userdata` — the pointer to pass to `callback` along with the
    ///   arguments when the closure is called
    ///
    /// # Result
    ///
    /// The new closure.
    pub fn new_mut<U, R>(cif:      Cif,
                         callback: CallbackMut<U, R>,
                         userdata: &'a mut U) -> Self
    {
        let cif = Box::new(cif);
        let (alloc, code) = low::closure_alloc();

        unsafe {
            low::prep_closure_mut(alloc,
                                  cif.as_raw_ptr(),
                                  callback,
                                  userdata as *mut U,
                                  code).unwrap();
        }

        Closure {
            _cif:    cif,
            alloc:   alloc,
            code:    code,
            _marker: PhantomData,
        }
    }

    /// Obtains the callable code pointer for a closure.
    ///
    /// # Safety
    ///
    /// The result needs to be transmuted to the correct type before
    /// it can be called. If the type is wrong then undefined behavior
    /// will result.
    pub fn code_ptr(&self) -> &unsafe extern "C" fn() {
        self.code.as_fun()
    }
}

#[cfg(test)]
mod test {
    use low;
    use super::*;
    use std::mem;
    use std::os::raw::c_void;

    #[test]
    fn call() {
        let cif  = Cif::new(vec![Type::i64(), Type::i64()].into_iter(),
                            Type::i64());
        let f    = |m: i64, n: i64| -> i64 {
            unsafe { cif.call(CodePtr(add_it as *mut c_void),
                              &[arg(&m), arg(&n)]) }
        };

        assert_eq!(12, f(5, 7));
        assert_eq!(13, f(6, 7));
        assert_eq!(15, f(8, 7));
    }

    extern "C" fn add_it(n: i64, m: i64) -> i64 {
        return n + m;
    }

    #[test]
    fn closure() {
        let cif  = Cif::new(vec![Type::u64()].into_iter(), Type::u64());
        let env: u64 = 5;
        let closure = Closure::new(cif, callback, &env);

        unsafe {
            let fun: &unsafe extern "C" fn(u64) -> u64
                = mem::transmute(closure.code_ptr());

            assert_eq!(11, fun(6));
            assert_eq!(12, fun(7));
        }
    }

    unsafe extern "C" fn callback(_cif: &low::ffi_cif,
                                  result: &mut u64,
                                  args: *const *const c_void,
                                  userdata: &u64)
    {
        let args: *const &u64 = mem::transmute(args);
        *result = **args + *userdata;
    }

    #[test]
    fn rust_lambda() {
        let cif = Cif::new(vec![Type::u64(), Type::u64()].into_iter(),
                           Type::u64());
        let env = |x: u64, y: u64| x + y;
        let closure = Closure::new(cif, callback2, &env);

        unsafe {
            let fun: &unsafe extern "C" fn (u64, u64) -> u64
                = mem::transmute(closure.code_ptr());

            assert_eq!(11, fun(5, 6));
        }
    }

    unsafe extern "C" fn callback2<F: Fn(u64, u64) -> u64>
        (_cif: &low::ffi_cif,
         result: &mut u64,
         args: *const *const c_void,
         userdata: &F)
    {
        let args: *const &u64 = mem::transmute(args);
        let arg1 = **args.offset(0);
        let arg2 = **args.offset(1);

        *result = userdata(arg1, arg2);
    }
}
<commit_msg>Fixed doc test that was not running.<commit_after>//! Middle layer providing a somewhat safer (but still quite unsafe)
//! API.
//!
//! The main idea of the middle layer is to wrap types
//! [`ffi_cif`](../raw/struct.ffi_cif.html) and
//! [`ffi_closure`](../raw/struct.ffi_closure.html) as
//! [`Cif`](struct.Cif.html) and [`Closure`](struct.Closure.html),
//! respectively, so that their resources are managed properly. However,
//! calling a function via a CIF or closure is still unsafe because
//! argument types aren’t checked. See the [`high`](../high/index.html)
//! layer for closures with type-checked arguments.
//!
use std::os::raw::c_void; use std::marker::PhantomData;

use low;
pub use low::{Callback, CallbackMut, CodePtr,
              ffi_abi as FfiAbi, FFI_DEFAULT_ABI};

pub mod types;
pub use self::types::Type;

/// Contains an untyped pointer to a function argument.
///
/// When calling a function via a [CIF](struct.Cif.html), each argument
/// must be passed as a C `void*`. Wrapping the argument in the `Arg`
/// struct accomplishes the necessary coercion.
#[derive(Debug)]
#[repr(C)]
pub struct Arg(*mut c_void);

impl Arg {
    /// Coerces an argument reference into the `Arg` type.
    ///
    /// This is used to wrap each argument pointer before passing them
    /// to [`Cif::call`](struct.Cif.html#method.call).
    pub fn new<T>(r: &T) -> Self {
        Arg(r as *const T as *mut c_void)
    }
}

/// Coerces an argument reference into the [`Arg`](struct.Arg.html)
/// type.
///
/// This is used to wrap each argument pointer before passing them
/// to [`Cif::call`](struct.Cif.html#method.call).
/// (This is the same as [`Arg::new`](struct.Arg.html#method.new)).
pub fn arg<T>(r: &T) -> Arg {
    Arg::new(r)
}

/// Describes the calling convention and types for calling a function.
///
/// This is the `middle` layer’s wrapping of the `low` and `raw` layers’
/// [`ffi_cif`](../raw/struct.ffi_cif.html). An initialized CIF contains
/// references to an array of argument types and a result type, each of
/// which may be allocated on the heap. `Cif` manages the memory of
/// those referenced objects.
///
/// Construct with [`Cif::new`](#method.new) or
/// [`Cif::from_type_array`](#method.from_type_array).
///
/// # Example
///
/// ```
/// extern "C" fn add(x: f64, y: &f64) -> f64 {
///     return x + y;
/// }
///
/// use libffi::middle::*;
///
/// let args = vec![Type::f64(), Type::pointer()];
/// let cif = Cif::new(args.into_iter(), Type::f64());
///
/// let n = unsafe { cif.call(CodePtr(add as _), &[arg(&5), arg(&&6)]) };
/// assert_eq!(11, n);
/// ```
#[derive(Clone, Debug)]
pub struct Cif {
    cif:    low::ffi_cif,
    args:   types::TypeArray,
    result: Type,
}

impl Cif {
    /// Creates a new CIF for the given argument and result types.
    ///
    /// Takes ownership of the argument and result
    /// [`Type`](types/struct.Type.html)s, because the resulting
    /// `Cif` retains references to them.
    /// Defaults to the platform’s default calling convention; this
    /// can be adjusted using [`set_abi`](#method.set_abi).
    pub fn new<I>(args: I, result: Type) -> Self
        where I: ExactSizeIterator<Item=Type>
    {
        Self::from_type_array(types::TypeArray::new(args), result)
    }

    /// Calls a function with the given arguments.
    ///
    /// In particular, this method invokes function `fun` passing it
    /// arguments `args`, and returns the result.
    ///
    /// # Safety
    ///
    /// There is no checking that the calling convention and types
    /// in the `Cif` match the actual calling convention and types of
    /// `fun`, nor that they match the types of `args`.
    pub unsafe fn call<R>(&self, fun: CodePtr, args: &[Arg]) -> R {
        use std::mem;

        assert!(self.cif.nargs as usize == args.len(),
                "Cif::call: passed wrong number of arguments");

        low::call::<R>(&self.cif as *const _ as *mut _,
                       fun,
                       mem::transmute::<*const Arg,
                                        *mut *mut c_void>(args.as_ptr()))
    }

    /// Sets the CIF to use the given calling convention.
    pub fn set_abi(&mut self, abi: FfiAbi) {
        self.cif.abi = abi;
    }

    /// Creates a new CIF for the given argument and result types.
    ///
    /// This is just like [`Cif::new`](#method.new), except it takes a
    /// `TypeArray` instead of an `ExactSizeIterator`.
    pub fn from_type_array(args: types::TypeArray, result: Type) -> Self {
        let mut cif: low::ffi_cif = Default::default();

        unsafe {
            low::prep_cif(&mut cif,
                          low::FFI_DEFAULT_ABI,
                          args.len(),
                          result.as_raw_ptr(),
                          args.as_raw_ptr())
        }.expect("low::prep_cif");

        // Note that cif retains references to args and result,
        // which is why we hold onto them here.
        Cif {
            cif:    cif,
            args:   args,
            result: result,
        }
    }

    /// Gets a raw pointer to the underlying
    /// [`ffi_cif`](../low/struct.ffi_cif.html).
    ///
    /// This can be used for passing a `middle::Cif` to functions from the
    /// [`low`](../low/index.html) and [`raw`](../raw/index.html) modules.
    pub fn as_raw_ptr(&self) -> *mut low::ffi_cif {
        &self.cif as *const _ as *mut _
    }
}

/// Represents a closure callable from C.
///
/// A libffi closure captures a `void*` (“userdata”) and passes it to a
/// callback when the code pointer (obtained via
/// [`code_ptr`](#method.code_ptr)) is invoked.
///
/// Construct with [`Closure::new`](#method.new) and
/// [`Closure::new_mut`](#method.new_mut).
///
/// # Example
///
/// In this example we turn a Rust lambda into a C function. Note that
/// the Rust closure value `lambda` is actually passed to the callback
/// function as userdata, which then invokes it.
///
/// ```
/// use std::mem;
/// use std::os::raw::c_void;
///
/// use libffi::middle::*;
/// use libffi::low;
///
/// unsafe extern "C" fn lambda_callback<F: Fn(u64, u64) -> u64>(
///     _cif: &low::ffi_cif,
///     result: &mut u64,
///     args: *const *const c_void,
///     userdata: &F)
/// {
///     let args: *const &u64 = mem::transmute(args);
///     let arg1 = **args.offset(0);
///     let arg2 = **args.offset(1);
///
///     *result = userdata(arg1, arg2);
/// }
///
/// let cif = Cif::new(vec![Type::u64(), Type::u64()].into_iter(),
///                    Type::u64());
/// let lambda = |x: u64, y: u64| x + y;
/// let closure = Closure::new(cif, lambda_callback, &lambda);
///
/// unsafe {
///     let fun: &unsafe extern "C" fn(u64, u64) -> u64
///         = mem::transmute(closure.code_ptr());
///
///     assert_eq!(11, fun(5, 6));
///     assert_eq!(12, fun(5, 7));
/// }
/// ```
#[derive(Debug)]
pub struct Closure<'a> {
    _cif:    Box<Cif>,
    alloc:   *mut ::low::ffi_closure,
    code:    CodePtr,
    _marker: PhantomData<&'a ()>,
}

impl<'a> Drop for Closure<'a> {
    fn drop(&mut self) {
        unsafe {
            low::closure_free(self.alloc);
        }
    }
}

impl<'a> Closure<'a> {
    /// Creates a new closure with immutable userdata.
    ///
    /// # Arguments
    ///
    /// - `cif` — describes the calling convention and argument and
    ///   result types
    /// - `callback` — the function to call when the closure is invoked
    /// - `userdata` — the pointer to pass to `callback` along with the
    ///   arguments when the closure is called
    ///
    /// # Result
    ///
    /// The new closure.
    pub fn new<U, R>(cif:      Cif,
                     callback: Callback<U, R>,
                     userdata: &'a U) -> Self
    {
        let cif = Box::new(cif);
        let (alloc, code) = low::closure_alloc();

        unsafe {
            low::prep_closure(alloc,
                              cif.as_raw_ptr(),
                              callback,
                              userdata as *const U,
                              code).unwrap();
        }

        Closure {
            _cif:    cif,
            alloc:   alloc,
            code:    code,
            _marker: PhantomData,
        }
    }

    /// Creates a new closure with mutable userdata.
    ///
    /// # Arguments
    ///
    /// - `cif` — describes the calling convention and argument and
    ///   result types
    /// - `callback` — the function to call when the closure is invoked
    /// - `userdata` — the pointer to pass to `callback` along with the
    ///   arguments when the closure is called
    ///
    /// # Result
    ///
    /// The new closure.
    pub fn new_mut<U, R>(cif:      Cif,
                         callback: CallbackMut<U, R>,
                         userdata: &'a mut U) -> Self
    {
        let cif = Box::new(cif);
        let (alloc, code) = low::closure_alloc();

        unsafe {
            low::prep_closure_mut(alloc,
                                  cif.as_raw_ptr(),
                                  callback,
                                  userdata as *mut U,
                                  code).unwrap();
        }

        Closure {
            _cif:    cif,
            alloc:   alloc,
            code:    code,
            _marker: PhantomData,
        }
    }

    /// Obtains the callable code pointer for a closure.
    ///
    /// # Safety
    ///
    /// The result needs to be transmuted to the correct type before
    /// it can be called. If the type is wrong then undefined behavior
    /// will result.
    pub fn code_ptr(&self) -> &unsafe extern "C" fn() {
        self.code.as_fun()
    }
}

#[cfg(test)]
mod test {
    use low;
    use super::*;
    use std::mem;
    use std::os::raw::c_void;

    #[test]
    fn call() {
        let cif  = Cif::new(vec![Type::i64(), Type::i64()].into_iter(),
                            Type::i64());
        let f    = |m: i64, n: i64| -> i64 {
            unsafe { cif.call(CodePtr(add_it as *mut c_void),
                              &[arg(&m), arg(&n)]) }
        };

        assert_eq!(12, f(5, 7));
        assert_eq!(13, f(6, 7));
        assert_eq!(15, f(8, 7));
    }

    extern "C" fn add_it(n: i64, m: i64) -> i64 {
        return n + m;
    }

    #[test]
    fn closure() {
        let cif  = Cif::new(vec![Type::u64()].into_iter(), Type::u64());
        let env: u64 = 5;
        let closure = Closure::new(cif, callback, &env);

        unsafe {
            let fun: &unsafe extern "C" fn(u64) -> u64
                = mem::transmute(closure.code_ptr());

            assert_eq!(11, fun(6));
            assert_eq!(12, fun(7));
        }
    }

    unsafe extern "C" fn callback(_cif: &low::ffi_cif,
                                  result: &mut u64,
                                  args: *const *const c_void,
                                  userdata: &u64)
    {
        let args: *const &u64 = mem::transmute(args);
        *result = **args + *userdata;
    }

    #[test]
    fn rust_lambda() {
        let cif = Cif::new(vec![Type::u64(), Type::u64()].into_iter(),
                           Type::u64());
        let env = |x: u64, y: u64| x + y;
        let closure = Closure::new(cif, callback2, &env);

        unsafe {
            let fun: &unsafe extern "C" fn (u64, u64) -> u64
                = mem::transmute(closure.code_ptr());

            assert_eq!(11, fun(5, 6));
        }
    }

    unsafe extern "C" fn callback2<F: Fn(u64, u64) -> u64>
        (_cif: &low::ffi_cif,
         result: &mut u64,
         args: *const *const c_void,
         userdata: &F)
    {
        let args: *const &u64 = mem::transmute(args);
        let arg1 = **args.offset(0);
        let arg2 = **args.offset(1);

        *result = userdata(arg1, arg2);
    }
}
<|endoftext|>


<commit_before>extern crate websocket_lowlevel;

use futures::future::Future;
use futures::stream::Stream;

use std::cell::RefCell;
use std::rc::Rc;

use super::{box_up_err, peer_err, peer_strerr, BoxedNewPeerFuture, Peer, Result};

use super::ws_peer::{Mode1, PeerForWs, WsReadWrapper, WsWriteWrapper};
use super::{once, ConstructParams, Options, PeerConstructor, Specifier};

use self::websocket_lowlevel::codec::ws::Context as WsLlContext;

#[derive(Debug, Clone)]
pub struct WsLlClient<T: Specifier>(pub T);
impl<T:Specifier> Specifier for WsLlClient<T> {
    fn construct(&self, p: ConstructParams) -> PeerConstructor {
        let inner = self.0.construct(p.clone());
        let opts = p.program_options;
        inner.map(move |q, _| get_ws_lowlevel_peer(
            WsLlContext::Client,
            q,
            opts.clone(),
        ))
    }
    specifier_boilerplate!(noglobalstate singleconnect has_subspec);
}
specifier_class!(
    name = WsLlClientClass,
    target = WsLlClient,
    prefixes = ["ws-lowlevel-client:","ws-ll-client:","ws-ll-c:"],
    arg_handling = subspec,
    overlay = false,
    MessageOriented,
    SingleConnect,
    help = r#"
[A] Low-level HTTP-independent WebSocket client connection without associated HTTP upgrade.

Example: TODO
"#
);

#[derive(Debug, Clone)]
pub struct WsLlServer<T: Specifier>(pub T);
impl<T:Specifier> Specifier for WsLlServer<T> {
    fn construct(&self, p: ConstructParams) -> PeerConstructor {
        let inner = self.0.construct(p.clone());
        let opts = p.program_options;
        inner.map(move |q, _| get_ws_lowlevel_peer(
            WsLlContext::Server,
            q,
            opts.clone(),
        ))
    }
    specifier_boilerplate!(noglobalstate singleconnect has_subspec);
}
specifier_class!(
    name = WsLlServerClass,
    target = WsLlServer,
    prefixes = ["ws-lowlevel-server:","ws-ll-server:","ws-ll-s:"],
    arg_handling = subspec,
    overlay = false,
    MessageOriented,
    SingleConnect,
    help = r#"
[A] Low-level HTTP-independent WebSocket server connection without associated HTTP upgrade.

Example: TODO
"#
);

pub fn get_ws_lowlevel_peer(mode: WsLlContext, inner: Peer, opts: Rc<Options>) -> BoxedNewPeerFuture {
    info!("get_ws_lowlevel_peer");
    unimplemented!()
}
<commit_msg>Implement ws-ll-{c,s}: specifier<commit_after>#![allow(unused)]

extern crate websocket_lowlevel;

use futures::future::Future;
use futures::stream::Stream;

use std::cell::RefCell;
use std::rc::Rc;

use super::{box_up_err, peer_err, peer_strerr, BoxedNewPeerFuture, Peer, Result};

use super::ws_peer::{Mode1, PeerForWs, WsReadWrapper, WsWriteWrapper};
use super::{once, ConstructParams, Options, PeerConstructor, Specifier};

use self::websocket_lowlevel::codec::ws::Context as WsLlContext;

#[derive(Debug, Clone)]
pub struct WsLlClient<T: Specifier>(pub T);
impl<T:Specifier> Specifier for WsLlClient<T> {
    fn construct(&self, p: ConstructParams) -> PeerConstructor {
        let inner = self.0.construct(p.clone());
        let opts = p.program_options;
        inner.map(move |q, _| get_ws_lowlevel_peer(
            WsLlContext::Client,
            q,
            opts.clone(),
        ))
    }
    specifier_boilerplate!(noglobalstate singleconnect has_subspec);
}
specifier_class!(
    name = WsLlClientClass,
    target = WsLlClient,
    prefixes = ["ws-lowlevel-client:","ws-ll-client:","ws-ll-c:"],
    arg_handling = subspec,
    overlay = false,
    MessageOriented,
    SingleConnect,
    help = r#"
[A] Low-level HTTP-independent WebSocket client connection without associated HTTP upgrade.

Example: TODO
"#
);

#[derive(Debug, Clone)]
pub struct WsLlServer<T: Specifier>(pub T);
impl<T:Specifier> Specifier for WsLlServer<T> {
    fn construct(&self, p: ConstructParams) -> PeerConstructor {
        let inner = self.0.construct(p.clone());
        let opts = p.program_options;
        inner.map(move |q, _| get_ws_lowlevel_peer(
            WsLlContext::Server,
            q,
            opts.clone(),
        ))
    }
    specifier_boilerplate!(noglobalstate singleconnect has_subspec);
}
specifier_class!(
    name = WsLlServerClass,
    target = WsLlServer,
    prefixes = ["ws-lowlevel-server:","ws-ll-server:","ws-ll-s:"],
    arg_handling = subspec,
    overlay = false,
    MessageOriented,
    SingleConnect,
    help = r#"
[A] Low-level HTTP-independent WebSocket server connection without associated HTTP upgrade.

Example: TODO
"#
);

pub fn get_ws_lowlevel_peer(mode: WsLlContext, inner: Peer, opts: Rc<Options>) -> BoxedNewPeerFuture {
    info!("get_ws_lowlevel_peer");

    use ::tokio_codec::Decoder;

    let c = websocket_lowlevel::codec::ws::MessageCodec::new(mode);
    let duplex = c.framed(PeerForWs(inner));

    let close_on_shutdown =  !opts.websocket_dont_close;
    let p = super::ws_peer::finish_building_ws_peer(&*opts, duplex, close_on_shutdown);

    Box::new(
        ::futures::future::ok(p)
    ) as BoxedNewPeerFuture
}
<|endoftext|>


<commit_before>use std::ops::{Add, Sub};

/// Basic point on a coordinate plane
///
/// The point contains two coordinates QR to describe its position in
/// two dimensions and a third T to describe its third dimension.
///
/// Q can also be considered X on a cube, and R is Z.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct Point(pub i32, pub i32, pub i32);

/// Add one point to another
impl<'a, 'b> Add<&'b Point> for &'a Point {
  type Output = Point;

  fn add(self, other: &'b Point) -> Point {
    let &Point(q0, r0, t0) = self;
    let &Point(q1, r1, t1) = other;

    Point(q0 + q1, r0 + r1, t0 + t1)
  }
}

/// Subtract one point from another
impl<'a, 'b> Sub<&'b Point> for &'a Point {
  type Output = Point;

  fn sub(self, other: &'b Point) -> Point {
    let &Point(q0, r0, t0) = self;
    let &Point(q1, r1, t1) = other;

    Point(q0 - q1, r0 - r1, t0 - t1)
  }
}

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

  #[test]
  fn add() {
    let point: Point = Point(1, 2, 5);
    let other: Point = Point(3, 4, 10);
    let Point(q, r, t) = &point + &other;

    assert!(4 == q);
    assert!(6 == r);
    assert!(15 == t);
  }

  #[test]
  fn sub() {
    let point: Point = Point(1, 2, 5);
    let other: Point = Point(3, 4, 10);
    let Point(q, r, t)= &point - &other;

    assert!(-2 == q);
    assert!(-2 == r);
    assert!(-5 == t);
  }
}
<commit_msg>Issue #86: Point is generic<commit_after>use std::ops::{Add, Sub};

/// Basic point on a coordinate plane
///
/// The point contains two coordinates QR to describe its position in
/// two dimensions and a third T to describe its third dimension.
///
/// Q can also be considered X on a cube, and R is Z.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct Point<T = i32>(pub T, pub T, pub T);

/// Add one point to another
impl<'a, 'b, T> Add<&'b Point<T>> for &'a Point<T>
  where &'a T: Add<&'b T, Output = T> {

  type Output = Point<T>;

  fn add(self, other: &'b Point<T>) -> Point<T> {
    let &Point(ref q0, ref r0, ref t0) = self;
    let &Point(ref q1, ref r1, ref t1) = other;

    Point(q0 + q1, r0 + r1, t0 + t1)
  }
}

/// Subtract one point from another
impl<'a, 'b, T> Sub<&'b Point<T>> for &'a Point<T>
  where &'a T: Sub<&'b T, Output = T> {

  type Output = Point<T>;

  fn sub(self, other: &'b Point<T>) -> Point<T> {
    let &Point(ref q0, ref r0, ref t0) = self;
    let &Point(ref q1, ref r1, ref t1) = other;

    Point(q0 - q1, r0 - r1, t0 - t1)
  }
}

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

  #[test]
  fn add() {
    let point: Point = Point(1, 2, 5);
    let other: Point = Point(3, 4, 10);
    let Point(q, r, t) = &point + &other;

    assert!(4 == q);
    assert!(6 == r);
    assert!(15 == t);
  }

  #[test]
  fn sub() {
    let point: Point = Point(1, 2, 5);
    let other: Point = Point(3, 4, 10);
    let Point(q, r, t)= &point - &other;

    assert!(-2 == q);
    assert!(-2 == r);
    assert!(-5 == t);
  }
}
<|endoftext|>


<commit_before>use {node, routing, time, hash};
use std::collections::VecDeque;
use std::str::FromStr;
use std::net;
use node::receptions;

pub const POLL_FREQUENCY_MS: u64 = 50;
pub const TRIES: u8 = 5;

#[test]
fn node_ping() {
   let alpha = node::Node::new().unwrap();
   let beta  = node::Node::new().unwrap();
   let beta_seed = beta.local_info();
   let span = time::Duration::seconds(1);

   // Bootstrapping alpha:
   assert!(alpha.bootstrap_until(beta_seed, 1).is_ok());

   let beta_receptions = alpha.receptions().during(span).from(beta.id().clone());

   // Alpha pings beta.
   assert!(alpha.ping(&beta.resources.id).is_ok());
   assert_eq!(1, beta_receptions.count());
}

#[test]
fn reception_iterator_times_out_correctly() {
   let alpha = node::Node::new().unwrap();
   let span = time::Duration::seconds(1);
   let maximum = time::Duration::seconds(3);
   let receptions = alpha.receptions().during(span);

   let before = time::SteadyTime::now();

   // nothing is happening, so this should time out in around a second (not necessarily precise)
   assert_eq!(0, receptions.count());

   let after = time::SteadyTime::now();

   assert!(after - before < maximum);
}

#[test]
fn bootstrapping_and_finding_on_simulated_network() {

   let mut nodes = simulated_network(30);

   // Head finds tail in a few steps.
   let head = nodes.pop_front().unwrap();
   let tail = nodes.pop_back().unwrap();

   assert_eq!(head.find_node(tail.id()).unwrap().id, tail.local_info().id);
}

#[test]
fn finding_on_simulated_unresponsive_network() {

   let mut nodes = simulated_network(30);
   nodes.drain(10..20);
   assert_eq!(nodes.len(), 20);

   // Head finds tail in a few steps.
   let head = nodes.pop_front().unwrap();
   let tail = nodes.pop_back().unwrap();

   assert_eq!(head.find_node(tail.id()).unwrap().id, tail.local_info().id);
}

#[test]
#[ignore]
fn finding_a_nonexisting_node_in_a_simulated_network_times_out() {

   let mut nodes = simulated_network(30);

   // Head finds tail in a few steps.
   let head = nodes.pop_front().unwrap();

   let random_hash = hash::SubotaiHash::random();
   assert!(head.find_node(&random_hash).is_err());
}

fn simulated_network(nodes: usize) -> VecDeque<node::Node> {
   let nodes: VecDeque<node::Node> = (0..nodes).map(|_| { node::Node::new().unwrap() }).collect();
   {
      let origin = node::Node::new().unwrap();

      // Initial handshake pass
      for node in nodes.iter() {
         assert!(node.bootstrap_until(origin.local_info(), 1).is_ok());
      }

      // Actual bootstrapping
      for node in nodes.iter() {
         assert!(node.bootstrap(origin.local_info()).is_ok());
      }
   }
   nodes
}

#[test]
fn updating_table_with_full_bucket_starts_the_conflict_resolution_mechanism()
{
   let node = node::Node::new().unwrap();
   node.resources.table.fill_bucket(8, routing::K as u8); // Bucket completely full

   let mut id = node.id().clone();
   id.flip_bit(8);
   id.raw[0] = 0xFF;
   let info = node_info_no_net(id);

   node.resources.update_table(info);
   assert_eq!(node.resources.conflicts.lock().unwrap().len(), 1);
}

#[test]
fn generating_a_conflict_causes_a_ping_to_the_evicted_node()
{
   let alpha = node::Node::new().unwrap();
   let beta = node::Node::new().unwrap();
   alpha.resources.update_table(beta.local_info());

   let index = alpha.resources.table.bucket_for_node(beta.id()).unwrap();

   // We fill the bucket corresponding to Beta until we are ready to cause a conflict.
   alpha.resources.table.fill_bucket(index, (routing::K -1) as u8);

   // We expect a ping to beta
   let pings = beta.receptions()
      .of_kind(receptions::KindFilter::Ping)
      .during(time::Duration::seconds(2));

   // Adding a new node causes a conflict.
   let mut id = beta.id().clone();
   id.raw[0] = 0xFF;
   let info = node_info_no_net(id);
   alpha.resources.update_table(info);

   assert_eq!(pings.count(), 1);
}

fn node_info_no_net(id : hash::SubotaiHash) -> routing::NodeInfo {
   routing::NodeInfo {
      id : id,
      address : net::SocketAddr::from_str("0.0.0.0:0").unwrap(),
   }
}
<commit_msg>Test for defensive state<commit_after>use {node, routing, time, hash};
use std::collections::VecDeque;
use std::str::FromStr;
use std::net;
use node::receptions;

pub const POLL_FREQUENCY_MS: u64 = 50;
pub const TRIES: u8 = 5;

#[test]
fn node_ping() {
   let alpha = node::Node::new().unwrap();
   let beta  = node::Node::new().unwrap();
   let beta_seed = beta.local_info();
   let span = time::Duration::seconds(1);

   // Bootstrapping alpha:
   assert!(alpha.bootstrap_until(beta_seed, 1).is_ok());

   let beta_receptions = alpha.receptions().during(span).from(beta.id().clone());

   // Alpha pings beta.
   assert!(alpha.ping(&beta.resources.id).is_ok());
   assert_eq!(1, beta_receptions.count());
}

#[test]
fn reception_iterator_times_out_correctly() {
   let alpha = node::Node::new().unwrap();
   let span = time::Duration::seconds(1);
   let maximum = time::Duration::seconds(3);
   let receptions = alpha.receptions().during(span);

   let before = time::SteadyTime::now();

   // nothing is happening, so this should time out in around a second (not necessarily precise)
   assert_eq!(0, receptions.count());

   let after = time::SteadyTime::now();

   assert!(after - before < maximum);
}

#[test]
fn bootstrapping_and_finding_on_simulated_network() {

   let mut nodes = simulated_network(30);

   // Head finds tail in a few steps.
   let head = nodes.pop_front().unwrap();
   let tail = nodes.pop_back().unwrap();

   assert_eq!(head.find_node(tail.id()).unwrap().id, tail.local_info().id);
}

#[test]
fn finding_on_simulated_unresponsive_network() {

   let mut nodes = simulated_network(30);
   nodes.drain(10..20);
   assert_eq!(nodes.len(), 20);

   // Head finds tail in a few steps.
   let head = nodes.pop_front().unwrap();
   let tail = nodes.pop_back().unwrap();

   assert_eq!(head.find_node(tail.id()).unwrap().id, tail.local_info().id);
}

#[test]
#[ignore]
fn finding_a_nonexisting_node_in_a_simulated_network_times_out() {

   let mut nodes = simulated_network(30);

   // Head finds tail in a few steps.
   let head = nodes.pop_front().unwrap();

   let random_hash = hash::SubotaiHash::random();
   assert!(head.find_node(&random_hash).is_err());
}

fn simulated_network(nodes: usize) -> VecDeque<node::Node> {
   let nodes: VecDeque<node::Node> = (0..nodes).map(|_| { node::Node::new().unwrap() }).collect();
   {
      let origin = node::Node::new().unwrap();

      // Initial handshake pass
      for node in nodes.iter() {
         assert!(node.bootstrap_until(origin.local_info(), 1).is_ok());
      }

      // Actual bootstrapping
      for node in nodes.iter() {
         assert!(node.bootstrap(origin.local_info()).is_ok());
      }
   }
   nodes
}

#[test]
fn updating_table_with_full_bucket_starts_the_conflict_resolution_mechanism()
{
   let node = node::Node::new().unwrap();
   node.resources.table.fill_bucket(8, routing::K as u8); // Bucket completely full

   let mut id = node.id().clone();
   id.flip_bit(8);
   id.raw[0] = 0xFF;
   let info = node_info_no_net(id);

   node.resources.update_table(info);
   assert_eq!(node.resources.conflicts.lock().unwrap().len(), 1);
}

#[test]
fn generating_a_conflict_causes_a_ping_to_the_evicted_node()
{
   let alpha = node::Node::new().unwrap();
   let beta = node::Node::new().unwrap();
   alpha.resources.update_table(beta.local_info());

   let index = alpha.resources.table.bucket_for_node(beta.id()).unwrap();

   // We fill the bucket corresponding to Beta until we are ready to cause a conflict.
   alpha.resources.table.fill_bucket(index, (routing::K -1) as u8);

   // We expect a ping to beta
   let pings = beta.receptions()
      .of_kind(receptions::KindFilter::Ping)
      .during(time::Duration::seconds(2));

   // Adding a new node causes a conflict.
   let mut id = beta.id().clone();
   id.raw[0] = 0xFF;
   let info = node_info_no_net(id);
   alpha.resources.update_table(info);

   assert_eq!(pings.count(), 1);
}

#[test]
fn generating_too_many_conflicts_causes_the_node_to_enter_defensive_state()
{
   let node = node::Node::new().unwrap();

   for index in 0..(routing::K + routing::MAX_CONFLICTS) {
      let mut id = node.id().clone();
      id.flip_bit(140); // Arbitrary bucket
      id.raw[0] = index as u8;
      let info = node_info_no_net(id);
      node.resources.update_table(info);
   }

   let state = *node.resources.state.read().unwrap();
   match state {
      node::State::Defensive => (),
      _ => panic!(),
   }
}

fn node_info_no_net(id : hash::SubotaiHash) -> routing::NodeInfo {
   routing::NodeInfo {
      id : id,
      address : net::SocketAddr::from_str("0.0.0.0:0").unwrap(),
   }
}
<|endoftext|>


<commit_before>use std::io::{self, Buffer, Reader, Writer, IoResult, BufferedStream, standard_error};
use std::mem;
use std::collections::BTreeMap;
use std::rand::{thread_rng, Rng};
use std::{u64, u16};
use std::num::{Int, FromPrimitive, ToPrimitive};
use url::Url;

#[cfg(test)]
use test::Bencher;
#[cfg(test)]
use serialize::json::ToJson;

use nonce::Nonce;
use message::{WSMessage, WSHeader, WS_MASK, WS_LEN, WS_LEN16, WS_LEN64, WS_OPCONT, WS_OPCODE, WS_OPTERM, WS_FIN};
use stream::NetworkStream;

pub struct WebSocket<S = NetworkStream> {
    stream: Option<BufferedStream<S>>,
    pub url: Url,
    hostname: String,
    use_ssl: bool,
}

impl WebSocket {
    pub fn new(url: Url) -> WebSocket {
        let use_ssl = url.scheme[] == "wss";

        let port = match url.port() {
            Some(p) => p,
            None if use_ssl => 443,
            _ => 80
        };

        WebSocket {
            stream: None,
            hostname: format!("{}:{}", url.serialize_host().unwrap(), port),
            url: url,
            use_ssl: use_ssl,
        }
    }

    fn try_connect(&mut self) -> IoResult<()> {
        self.stream = Some(BufferedStream::new(try!(NetworkStream::connect(self.hostname[], self.use_ssl))));
        Ok(())
    }

    fn write_request(&mut self, nonce: &str) -> IoResult<()> {
        let s = match self.stream { Some(ref mut s) => s, None => return Err(standard_error(io::IoErrorKind::NotConnected)) };

        try!(s.write(format!("GET {} HTTP/1.1\r\n", self.url.serialize_path().unwrap_or("/".to_string())).as_bytes()));
        try!(s.write(format!("Host: {}\r\n", self.url.host().unwrap()).as_bytes()));
        try!(s.write(format!("Origin: {}\r\n", self.url.serialize_no_fragment()).as_bytes()));
        try!(s.write(format!("Sec-WebSocket-Key: {}\r\n", nonce).as_bytes()));

        try!(s.write(b"Upgrade: websocket\r\n"));
        try!(s.write(b"Connection: Upgrade\r\n"));
        try!(s.write(b"Sec-WebSocket-Protocol: char, superchat\r\n"));
        try!(s.write(b"Sec-WebSocket-Version: 13\r\n"));
        try!(s.write(b"\r\n"));

        s.flush()
    }

    fn read_response(&mut self, nonce: &str) -> IoResult<()> {
        let spaces: &[_] = &[' ', '\t', '\r', '\n'];
        let s = match self.stream { Some(ref mut s) => s, None => return Err(standard_error(io::IoErrorKind::NotConnected)) };
        let status = try!(s.read_line())[].splitn(2, ' ').nth(1).and_then(|s| s.parse::<uint>());

        match status {
            Some(101) => (),
            _ => return Err(standard_error(io::InvalidInput))
        }

        let headers = s.lines().map(|r| r.unwrap_or("\r\n".to_string())) .take_while(|l| l[] != "\r\n")
            .map(|s| s[].splitn(1, ':').map(|s| s.trim_matches(spaces).to_string()).collect::<Vec<String>>())
            .map(|p| (p[0].to_string(), p[1].to_string()))
            .collect::<BTreeMap<String, String>>();

        try!(s.flush());

        let response = headers.get("Sec-WebSocket-Accept");
        match response {
            Some(r) if nonce == r[] => (),
            _ => return Err(standard_error(io::InvalidInput))
        }

        Ok(())
    }

    pub fn connect(&mut self) -> IoResult<()> {
        let mut nonce = Nonce::new();

        try!(self.try_connect());
        try!(self.write_request(nonce[]));

        nonce = nonce.encode();
        try!(self.read_response(nonce[]));

        Ok(())
    }

    fn read_header(&mut self) -> IoResult<WSHeader> {
        // XXX: this is a bug, WSHeader should accept u16
        Ok(WSHeader::from_bits_truncate(try!(self.read_be_u16())))
    }

    fn read_length(&mut self, header: &WSHeader) -> IoResult<uint> {
        let wslen = *header & WS_LEN;
        if wslen == WS_LEN16 { self.read_be_u16().map(|v| v as uint) }
        else if wslen == WS_LEN64 { self.read_be_u64().map(|v| v as uint) }
        else { Ok(wslen.bits() as uint) }
    }

    pub fn read_message(&mut self) -> IoResult<WSMessage> {
        let header = try!(self.read_header());
        let mut len = try!(self.read_length(&header));

        let mask = if header.contains(WS_MASK) {
            Some(try!(self.read_be_u32()))
        } else {
            None
        };

        let mut status = if header.contains(WS_OPTERM) {
            len = len - 2;
            Some(try!(self.read_be_u16()))
        } else {
            None
        };

        let mut data = try!(self.read_exact(len));
        if let Some(mut m) = mask {
            if let Some(s) = status {
                status = Some(s ^ (m & 0xffff) as u16);
                m = m.rotate_right(16);
            }
            data = WebSocket::mask_data(data[], m);
        }

        Ok(WSMessage { header: header, data: data, status: status.and_then(FromPrimitive::from_u16) })
    }

    fn mask_data(data: &[u8], mask: u32) -> Vec<u8> {
        data.iter().enumerate().map(|(i, b)| *b ^ (mask >> ((i % 4) << 3) & 0xff) as u8).collect::<Vec<u8>>()
    }

    pub fn send_message(&mut self, msg: &WSMessage) -> IoResult<()> {
        let mut len = msg.data.len();
        let mut hdr = msg.header - WS_LEN;

        if msg.status.is_some() {
            len = len + 2;
        }

        if len < u16::MAX as uint {
            hdr = hdr | WSHeader::from_bits_truncate(len as u16);
            try!(self.write_be_u16(hdr.bits()));

        } else if len < u64::MAX as uint {
            hdr = hdr | WS_LEN16;
            try!(self.write_be_u16(hdr.bits()));
            try!(self.write_be_u16(len as u16));

        } else {
            hdr = hdr | WS_LEN64;
            try!(self.write_be_u16(hdr.bits()));
            try!(self.write_be_u64(len as u64));
        }

        if hdr.contains(WS_MASK) {
            let mut mask = thread_rng().gen::<u32>();
            try!(self.write_be_u32(mask));

            if let Some(status) = msg.status {
                try!(self.write_be_u16(status.to_u16().unwrap() ^ (mask & 0xffff) as u16));
                mask = mask.rotate_right(16);
            }

            try!(self.write(WebSocket::mask_data(msg.data[], mask)[]));
        } else {
            if let Some(status) = msg.status {
                try!(self.write_be_u16(status.to_u16().unwrap()));
            }
            try!(self.write(msg.data[]));
        }

        self.flush()
    }

    pub fn iter(&mut self) -> WSMessages {
        WSMessages { sock: self }
    }
}

impl Reader for WebSocket {
    fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
        match self.stream {
            Some(ref mut s) => s.read(buf),
            None => Err(standard_error(io::IoErrorKind::NotConnected))
        }
    }
}

impl Writer for WebSocket {
    fn write(&mut self, buf: &[u8]) -> IoResult<()> {
        match self.stream {
            Some(ref mut s) => s.write(buf),
            None => Err(standard_error(io::IoErrorKind::NotConnected))
        }
    }

    fn flush(&mut self) -> IoResult<()> {
        match self.stream {
            Some(ref mut s) => s.flush(),
            None => Err(standard_error(io::IoErrorKind::NotConnected))
        }
    }
}

impl Buffer for WebSocket {
    fn fill_buf<'a>(&'a mut self) -> IoResult<&'a [u8]> {
        match self.stream {
            Some(ref mut s) => s.fill_buf(),
            None => Err(standard_error(io::IoErrorKind::NotConnected))
        }
    }

    fn consume(&mut self, amt: uint) {
        match self.stream {
            Some(ref mut s) => s.consume(amt),
            None => ()
        }
    }
}

pub struct WSMessages<'a> {
    sock: &'a mut WebSocket
}

pub struct WSDefragMessages<'a> {
    underlying: &'a mut WSMessages<'a>,
    buffer: WSMessage
}

impl<'a> WSMessages<'a> {
    pub fn defrag(&'a mut self) -> WSDefragMessages<'a> {
        WSDefragMessages{ underlying: self, buffer: WSMessage{ header: WSHeader::empty(), data: Vec::new(), status: None } }
    }
}

impl<'a> Iterator for WSMessages<'a> {
    type Item = WSMessage;
    fn next(&mut self) -> Option<WSMessage> {
        self.sock.read_message().ok()
    }
}

impl<'a> WSDefragMessages<'a> {
    fn popbuf(&mut self) -> Option<WSMessage> {
        if self.buffer.data.is_empty() {
            None
        } else {
            let mut buf = WSMessage{ header: WSHeader::empty(), data: Vec::new(), status: None };
            mem::swap(&mut self.buffer, &mut buf);
            Some(buf)
        }
    }

    fn swapbuf(&mut self, msg: &mut WSMessage) -> () {
        mem::swap(&mut self.buffer, msg);
    }
}

impl<'a> Iterator for WSDefragMessages<'a> {
    type Item = WSMessage;
    fn next(&mut self) -> Option<WSMessage> {
        loop {
            match self.underlying.next() {
                None => return self.popbuf(),
                Some(mut msg) => if msg.header.contains(WS_FIN) {
                    if msg.header & WS_OPCODE == WS_OPCONT {
                        self.buffer.push(msg);
                        return self.popbuf();
                    } else {
                        return Some(msg);
                    }

                } else {
                    if msg.header & WS_OPCODE == WS_OPCONT {
                        self.buffer.push(msg);
                    } else {
                        self.swapbuf(&mut msg);
                        return Some(msg);
                    }
                }
            }
        }
    }
}

#[bench]
#[allow(dead_code, unused_variables)]
fn test_connect(b: &mut Bencher) {
    //let url = Url::parse("wss://stream.pushbullet.com/websocket/").unwrap();
    let url = Url::parse("ws://echo.websocket.org").unwrap();
    let mut ws = WebSocket::new(url);
    ws.connect().unwrap();

    let msg = WSMessage::text("Hello, World!");
    b.bytes = msg.data.len() as u64;

    b.auto_bench(|b| {
        b.iter(|| {
            ws.send_message(&msg).unwrap();
            let msg = ws.read_message().unwrap();
            //println!("received: {} {}", msg, msg.to_string());
        })
    });

    //for msg in ws.iter() {
        //println!("{}", msg.to_string());
    //}
}
<commit_msg>add some comments<commit_after>use std::io::{self, Buffer, Reader, Writer, IoResult, BufferedStream, standard_error};
use std::mem;
use std::collections::BTreeMap;
use std::rand::{thread_rng, Rng};
use std::{u64, u16};
use std::num::{Int, FromPrimitive, ToPrimitive};
use url::Url;

#[cfg(test)]
use test::Bencher;
#[cfg(test)]
use serialize::json::ToJson;

use nonce::Nonce;
use message::{WSMessage, WSHeader, WS_MASK, WS_LEN, WS_LEN16, WS_LEN64, WS_OPCONT, WS_OPCODE, WS_OPTERM, WS_FIN};
use stream::NetworkStream;

pub struct WebSocket<S = NetworkStream> {
    stream: Option<BufferedStream<S>>,
    pub url: Url,
    hostname: String,
    use_ssl: bool,
}

impl WebSocket {
    pub fn new(url: Url) -> WebSocket {
        let use_ssl = url.scheme[] == "wss";

        let port = match url.port() {
            Some(p) => p,
            None if use_ssl => 443,
            _ => 80
        };

        WebSocket {
            stream: None,
            hostname: format!("{}:{}", url.serialize_host().unwrap(), port),
            url: url,
            use_ssl: use_ssl,
        }
    }

    fn try_connect(&mut self) -> IoResult<()> {
        self.stream = Some(BufferedStream::new(try!(NetworkStream::connect(self.hostname[], self.use_ssl))));
        Ok(())
    }

    fn write_request(&mut self, nonce: &str) -> IoResult<()> {
        let s = match self.stream { Some(ref mut s) => s, None => return Err(standard_error(io::IoErrorKind::NotConnected)) };

        try!(s.write(format!("GET {} HTTP/1.1\r\n", self.url.serialize_path().unwrap_or("/".to_string())).as_bytes()));
        try!(s.write(format!("Host: {}\r\n", self.url.host().unwrap()).as_bytes()));
        try!(s.write(format!("Origin: {}\r\n", self.url.serialize_no_fragment()).as_bytes()));
        try!(s.write(format!("Sec-WebSocket-Key: {}\r\n", nonce).as_bytes()));

        try!(s.write(b"Upgrade: websocket\r\n"));
        try!(s.write(b"Connection: Upgrade\r\n"));
        try!(s.write(b"Sec-WebSocket-Protocol: char, superchat\r\n"));
        try!(s.write(b"Sec-WebSocket-Version: 13\r\n"));
        try!(s.write(b"\r\n"));

        s.flush()
    }

    fn read_response(&mut self, nonce: &str) -> IoResult<()> {
        let spaces: &[_] = &[' ', '\t', '\r', '\n'];
        let s = match self.stream { Some(ref mut s) => s, None => return Err(standard_error(io::IoErrorKind::NotConnected)) };
        let status = try!(s.read_line())[].splitn(2, ' ').nth(1).and_then(|s| s.parse::<uint>());

        match status {
            Some(101) => (),
            _ => return Err(standard_error(io::InvalidInput))
        }

        let headers = s.lines().map(|r| r.unwrap_or("\r\n".to_string())) .take_while(|l| l[] != "\r\n")
            .map(|s| s[].splitn(1, ':').map(|s| s.trim_matches(spaces).to_string()).collect::<Vec<String>>())
            .map(|p| (p[0].to_string(), p[1].to_string()))
            .collect::<BTreeMap<String, String>>();

        try!(s.flush());

        let response = headers.get("Sec-WebSocket-Accept");
        match response {
            Some(r) if nonce == r[] => (),
            _ => return Err(standard_error(io::InvalidInput))
        }

        Ok(())
    }

    pub fn connect(&mut self) -> IoResult<()> {
        let mut nonce = Nonce::new();

        try!(self.try_connect());
        try!(self.write_request(nonce[]));

        nonce = nonce.encode();
        try!(self.read_response(nonce[]));

        Ok(())
    }

    fn read_header(&mut self) -> IoResult<WSHeader> {
        // XXX: this is a bug, WSHeader should accept u16
        Ok(WSHeader::from_bits_truncate(try!(self.read_be_u16())))
    }

    fn read_length(&mut self, header: &WSHeader) -> IoResult<uint> {
        let wslen = *header & WS_LEN;
        if wslen == WS_LEN16 { self.read_be_u16().map(|v| v as uint) }
        else if wslen == WS_LEN64 { self.read_be_u64().map(|v| v as uint) }
        else { Ok(wslen.bits() as uint) }
    }

    pub fn read_message(&mut self) -> IoResult<WSMessage> {
        let header = try!(self.read_header());
        let mut len = try!(self.read_length(&header));

        let mask = if header.contains(WS_MASK) {
            Some(try!(self.read_be_u32()))
        } else {
            None
        };

        // If this is the terminating frame (close command),
        // first two bytes of data MUST BE u16 status code
        let mut status = if header.contains(WS_OPTERM) {
            // compensate length of status code
            len = len - 2;
            Some(try!(self.read_be_u16()))
        } else {
            None
        };

        let mut data = try!(self.read_exact(len));

        // If we have mask, decrypt data
        if let Some(mut m) = mask {
            // decrypt status if present
            if let Some(s) = status {
                status = Some(s ^ (m & 0xffff) as u16);
                // compensate the usage of two mask bytes
                m = m.rotate_right(16);
            }
            data = WebSocket::mask_data(data[], m);
        }

        Ok(WSMessage { header: header, data: data, status: status.and_then(FromPrimitive::from_u16) })
    }

    fn mask_data(data: &[u8], mask: u32) -> Vec<u8> {
        data.iter().enumerate().map(|(i, b)| *b ^ (mask >> ((i % 4) << 3) & 0xff) as u8).collect::<Vec<u8>>()
    }

    pub fn send_message(&mut self, msg: &WSMessage) -> IoResult<()> {
        let mut len = msg.data.len();
        let mut hdr = msg.header - WS_LEN;

        // If we have status set, the data length is increased by status size
        if msg.status.is_some() {
            len = len + 2;
        }

        // Encode and send length along with header
        if len < u16::MAX as uint {
            hdr = hdr | WSHeader::from_bits_truncate(len as u16);
            try!(self.write_be_u16(hdr.bits()));

        } else if len < u64::MAX as uint {
            hdr = hdr | WS_LEN16;
            try!(self.write_be_u16(hdr.bits()));
            try!(self.write_be_u16(len as u16));

        } else {
            hdr = hdr | WS_LEN64;
            try!(self.write_be_u16(hdr.bits()));
            try!(self.write_be_u64(len as u64));
        }

        // If user required masking, encrypt all data
        if hdr.contains(WS_MASK) {
            // Generate and send random mask
            let mut mask = thread_rng().gen::<u32>();
            try!(self.write_be_u32(mask));

            // Encrypt status code if present
            if let Some(status) = msg.status {
                try!(self.write_be_u16(status.to_u16().unwrap() ^ (mask & 0xffff) as u16));
                // compensate for mask already used for status encryption
                mask = mask.rotate_right(16);
            }

            try!(self.write(WebSocket::mask_data(msg.data[], mask)[]));
        } else {
            // Send status code if present
            if let Some(status) = msg.status {
                try!(self.write_be_u16(status.to_u16().unwrap()));
            }
            try!(self.write(msg.data[]));
        }

        self.flush()
    }

    pub fn iter(&mut self) -> WSMessages {
        WSMessages { sock: self }
    }
}

impl Reader for WebSocket {
    fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
        match self.stream {
            Some(ref mut s) => s.read(buf),
            None => Err(standard_error(io::IoErrorKind::NotConnected))
        }
    }
}

impl Writer for WebSocket {
    fn write(&mut self, buf: &[u8]) -> IoResult<()> {
        match self.stream {
            Some(ref mut s) => s.write(buf),
            None => Err(standard_error(io::IoErrorKind::NotConnected))
        }
    }

    fn flush(&mut self) -> IoResult<()> {
        match self.stream {
            Some(ref mut s) => s.flush(),
            None => Err(standard_error(io::IoErrorKind::NotConnected))
        }
    }
}

impl Buffer for WebSocket {
    fn fill_buf<'a>(&'a mut self) -> IoResult<&'a [u8]> {
        match self.stream {
            Some(ref mut s) => s.fill_buf(),
            None => Err(standard_error(io::IoErrorKind::NotConnected))
        }
    }

    fn consume(&mut self, amt: uint) {
        match self.stream {
            Some(ref mut s) => s.consume(amt),
            None => ()
        }
    }
}

pub struct WSMessages<'a> {
    sock: &'a mut WebSocket
}

pub struct WSDefragMessages<'a> {
    underlying: &'a mut WSMessages<'a>,
    buffer: WSMessage
}

impl<'a> WSMessages<'a> {
    pub fn defrag(&'a mut self) -> WSDefragMessages<'a> {
        WSDefragMessages{ underlying: self, buffer: WSMessage{ header: WSHeader::empty(), data: Vec::new(), status: None } }
    }
}

impl<'a> Iterator for WSMessages<'a> {
    type Item = WSMessage;
    fn next(&mut self) -> Option<WSMessage> {
        self.sock.read_message().ok()
    }
}

impl<'a> WSDefragMessages<'a> {
    fn popbuf(&mut self) -> Option<WSMessage> {
        if self.buffer.data.is_empty() {
            None
        } else {
            let mut buf = WSMessage{ header: WSHeader::empty(), data: Vec::new(), status: None };
            mem::swap(&mut self.buffer, &mut buf);
            Some(buf)
        }
    }

    fn swapbuf(&mut self, msg: &mut WSMessage) -> () {
        mem::swap(&mut self.buffer, msg);
    }
}

impl<'a> Iterator for WSDefragMessages<'a> {
    type Item = WSMessage;
    fn next(&mut self) -> Option<WSMessage> {
        loop {
            match self.underlying.next() {
                None => return self.popbuf(),
                Some(mut msg) => if msg.header.contains(WS_FIN) {
                    if msg.header & WS_OPCODE == WS_OPCONT {
                        self.buffer.push(msg);
                        return self.popbuf();
                    } else {
                        return Some(msg);
                    }

                } else {
                    if msg.header & WS_OPCODE == WS_OPCONT {
                        self.buffer.push(msg);
                    } else {
                        self.swapbuf(&mut msg);
                        return Some(msg);
                    }
                }
            }
        }
    }
}

#[bench]
#[allow(dead_code, unused_variables)]
fn test_connect(b: &mut Bencher) {
    //let url = Url::parse("wss://stream.pushbullet.com/websocket/").unwrap();
    let url = Url::parse("ws://echo.websocket.org").unwrap();
    let mut ws = WebSocket::new(url);
    ws.connect().unwrap();

    let msg = WSMessage::text("Hello, World!");
    b.bytes = msg.data.len() as u64;

    b.auto_bench(|b| {
        b.iter(|| {
            ws.send_message(&msg).unwrap();
            let msg = ws.read_message().unwrap();
            //println!("received: {} {}", msg, msg.to_string());
        })
    });

    //for msg in ws.iter() {
        //println!("{}", msg.to_string());
    //}
}
<|endoftext|>