From c620411d9222cfcb287b2012b7f53f0d8877912c Mon Sep 17 00:00:00 2001 From: Sergio Benitez Date: Sat, 17 Feb 2018 04:06:05 -0800 Subject: [PATCH] Add 'keep_alive' configuration parameter. The 'keep_alive' configuration parameter allows HTTP keep-alive timeouts to be configured or for keep-alive to be disabled entirely. --- examples/config/Rocket.toml | 5 +- lib/src/config/builder.rs | 31 ++++++++++++ lib/src/config/config.rs | 40 ++++++++++++++-- lib/src/config/custom_values.rs | 19 ++++++++ lib/src/config/mod.rs | 84 +++++++++++++++++++++++++++++++++ lib/src/config/toml_ext.rs | 1 + lib/src/rocket.rs | 14 +++++- 7 files changed, 189 insertions(+), 5 deletions(-) diff --git a/examples/config/Rocket.toml b/examples/config/Rocket.toml index 18befb37..e1e2b673 100644 --- a/examples/config/Rocket.toml +++ b/examples/config/Rocket.toml @@ -10,6 +10,7 @@ msgpack = 1048576 # this is an extra used by the msgpack contrib module address = "localhost" port = 8000 workers = 1 +keep_alive = 5 log = "normal" hi = "Hello!" # this is an unused extra; maybe application specific? is_extra = true # this is an unused extra; maybe application specific? @@ -17,8 +18,9 @@ is_extra = true # this is an unused extra; maybe application specific? [staging] address = "0.0.0.0" port = 8000 -log = "normal" workers = 8 +keep_alive = 5 +log = "normal" # don't use this key! generate your own and keep it private! secret_key = "8Xui8SN4mI+7egV/9dlfYYLGQJeEx4+DwmSQLwDVXJg=" @@ -26,6 +28,7 @@ secret_key = "8Xui8SN4mI+7egV/9dlfYYLGQJeEx4+DwmSQLwDVXJg=" address = "0.0.0.0" port = 8000 workers = 12 +keep_alive = 5 log = "critical" # don't use this key! generate your own and keep it private! secret_key = "hPRYyVRiMyxpw5sBB1XeCMN1kFsDCqKvBi2QJxBVHQk=" diff --git a/lib/src/config/builder.rs b/lib/src/config/builder.rs index 8eaaf114..7a8b3f0e 100644 --- a/lib/src/config/builder.rs +++ b/lib/src/config/builder.rs @@ -14,6 +14,8 @@ pub struct ConfigBuilder { pub port: u16, /// The number of workers to run in parallel. pub workers: u16, + /// Keep-alive timeout in seconds or None if disabled. + pub keep_alive: Option, /// How much information to log. pub log_level: LoggingLevel, /// The secret key. @@ -63,6 +65,7 @@ impl ConfigBuilder { address: config.address, port: config.port, workers: config.workers, + keep_alive: config.keep_alive, log_level: config.log_level, secret_key: None, tls: None, @@ -128,6 +131,32 @@ impl ConfigBuilder { self } + /// Set the keep-alive timeout to `timeout` seconds. If `timeout` is `None`, + /// keep-alive is disabled. + /// + /// # Example + /// + /// ```rust + /// use rocket::config::{Config, Environment}; + /// + /// let config = Config::build(Environment::Staging) + /// .keep_alive(10) + /// .unwrap(); + /// + /// assert_eq!(config.keep_alive, Some(10)); + /// + /// let config = Config::build(Environment::Staging) + /// .keep_alive(None) + /// .unwrap(); + /// + /// assert_eq!(config.keep_alive, None); + /// ``` + #[inline] + pub fn keep_alive>>(mut self, timeout: T) -> Self { + self.keep_alive = timeout.into(); + self + } + /// Sets the `log_level` in the configuration being built. /// /// # Example @@ -283,6 +312,7 @@ impl ConfigBuilder { /// .address("127.0.0.1") /// .port(700) /// .workers(12) + /// .keep_alive(None) /// .finalize(); /// /// assert!(config.is_ok()); @@ -298,6 +328,7 @@ impl ConfigBuilder { config.set_address(self.address)?; config.set_port(self.port); config.set_workers(self.workers); + config.set_keep_alive(self.keep_alive); config.set_log_level(self.log_level); config.set_extras(self.extras); config.set_root(self.root); diff --git a/lib/src/config/config.rs b/lib/src/config/config.rs index af60317d..31c18536 100644 --- a/lib/src/config/config.rs +++ b/lib/src/config/config.rs @@ -47,6 +47,8 @@ pub struct Config { pub port: u16, /// The number of workers to run concurrently. pub workers: u16, + /// Keep-alive timeout in seconds or None if disabled. + pub keep_alive: Option, /// How much information to log. pub log_level: LoggingLevel, /// The secret key. @@ -63,7 +65,7 @@ pub struct Config { macro_rules! config_from_raw { ($config:expr, $name:expr, $value:expr, - $($key:ident => ($type:ident, $set:ident, $map:expr)),+ | _ => $rest:expr) => ( + $($key:ident => ($type:ident, $set:ident, $map:expr),)+ | _ => $rest:expr) => ( match $name { $(stringify!($key) => { super::custom_values::$type($config, $name, $value) @@ -213,6 +215,7 @@ impl Config { address: "localhost".to_string(), port: 8000, workers: default_workers, + keep_alive: Some(5), log_level: LoggingLevel::Normal, secret_key: key, tls: None, @@ -227,6 +230,7 @@ impl Config { address: "0.0.0.0".to_string(), port: 8000, workers: default_workers, + keep_alive: Some(5), log_level: LoggingLevel::Normal, secret_key: key, tls: None, @@ -241,6 +245,7 @@ impl Config { address: "0.0.0.0".to_string(), port: 8000, workers: default_workers, + keep_alive: Some(5), log_level: LoggingLevel::Critical, secret_key: key, tls: None, @@ -275,6 +280,7 @@ impl Config { /// * **address**: String /// * **port**: Integer (16-bit unsigned) /// * **workers**: Integer (16-bit unsigned) + /// * **keep_alive**: Integer or Boolean (false) or String ('none') /// * **log**: String /// * **secret_key**: String (192-bit base64) /// * **tls**: Table (`certs` (path as String), `key` (path as String)) @@ -284,10 +290,11 @@ impl Config { address => (str, set_address, id), port => (u16, set_port, ok), workers => (u16, set_workers, ok), - secret_key => (str, set_secret_key, id), + keep_alive => (u32_option, set_keep_alive, ok), log => (log_level, set_log_level, ok), + secret_key => (str, set_secret_key, id), tls => (tls_config, set_raw_tls, id), - limits => (limits, set_limits, ok) + limits => (limits, set_limits, ok), | _ => { self.extras.insert(name.into(), val.clone()); Ok(()) @@ -390,6 +397,31 @@ impl Config { self.workers = workers; } + /// Set the keep-alive timeout to `timeout` seconds. If `timeout` is `None`, + /// keep-alive is disabled. + /// + /// # Example + /// + /// ```rust + /// use rocket::config::Config; + /// + /// # use rocket::config::ConfigError; + /// # fn config_test() -> Result<(), ConfigError> { + /// let mut config = Config::development()?; + /// + /// // Set keep-alive timeout to 10 seconds. + /// config.set_keep_alive(10); + /// + /// // Disable keep-alive. + /// config.set_keep_alive(None); + /// # Ok(()) + /// # } + /// ``` + #[inline] + pub fn set_keep_alive>>(&mut self, timeout: T) { + self.keep_alive = timeout.into(); + } + /// Sets the `secret_key` in `self` to `key` which must be a 192-bit base64 /// encoded string. /// @@ -853,6 +885,7 @@ impl fmt::Debug for Config { s.field("address", &self.address); s.field("port", &self.port); s.field("workers", &self.workers); + s.field("keep_alive", &self.keep_alive); s.field("log_level", &self.log_level); for (key, value) in self.extras() { @@ -870,6 +903,7 @@ impl PartialEq for Config { && self.port == other.port && self.workers == other.workers && self.log_level == other.log_level + && self.keep_alive == other.keep_alive && self.environment == other.environment && self.extras == other.extras } diff --git a/lib/src/config/custom_values.rs b/lib/src/config/custom_values.rs index bd739a26..2dce6d6f 100644 --- a/lib/src/config/custom_values.rs +++ b/lib/src/config/custom_values.rs @@ -3,6 +3,7 @@ use std::fmt; #[cfg(feature = "tls")] use rustls::{Certificate, PrivateKey}; use config::{Result, Config, Value, ConfigError, LoggingLevel}; +use http::uncased::uncased_eq; use http::Key; #[derive(Clone)] @@ -260,3 +261,21 @@ pub fn limits(conf: &Config, name: &str, value: &Value) -> Result { Ok(limits) } + +pub fn u32_option(conf: &Config, name: &str, value: &Value) -> Result> { + let expect = "a 32-bit unsigned integer or 'none' or 'false'"; + let err = Err(conf.bad_type(name, value.type_str(), expect)); + + match value.as_integer() { + Some(x) if x >= 0 && x <= (u32::max_value() as i64) => Ok(Some(x as u32)), + Some(_) => err, + None => match value.as_str() { + Some(v) if uncased_eq(v, "none") => Ok(None), + Some(_) => err, + _ => match value.as_bool() { + Some(false) => Ok(None), + _ => err + } + } + } +} diff --git a/lib/src/config/mod.rs b/lib/src/config/mod.rs index fee43999..313a773a 100644 --- a/lib/src/config/mod.rs +++ b/lib/src/config/mod.rs @@ -38,6 +38,9 @@ //! * examples: `"8000"`, `"80"`, `"4242"` //! * **workers**: _[integer]_ the number of concurrent workers to use //! * examples: `12`, `1`, `4` +//! * **keep_alive**: _[integer, 'false', or 'none']_ timeout, in seconds, for +//! HTTP keep-alive. disabled on 'false' or 'none' +//! * examples: `5`, `60`, `false`, `"none"` //! * **log**: _[string]_ how much information to log; one of `"normal"`, //! `"debug"`, or `"critical"` //! * **secret_key**: _[string]_ a 256-bit base64 encoded string (44 @@ -72,6 +75,7 @@ //! address = "localhost" //! port = 8000 //! workers = [number_of_cpus * 2] +//! keep_alive = 5 //! log = "normal" //! secret_key = [randomly generated at launch] //! limits = { forms = 32768 } @@ -80,6 +84,7 @@ //! address = "0.0.0.0" //! port = 8000 //! workers = [number_of_cpus * 2] +//! keep_alive = 5 //! log = "normal" //! secret_key = [randomly generated at launch] //! limits = { forms = 32768 } @@ -88,6 +93,7 @@ //! address = "0.0.0.0" //! port = 8000 //! workers = [number_of_cpus * 2] +//! keep_alive = 5 //! log = "critical" //! secret_key = [randomly generated at launch] //! limits = { forms = 32768 } @@ -601,6 +607,7 @@ mod test { port = 7810 workers = 21 log = "critical" + keep_alive = false secret_key = "8Xui8SN4mI+7egV/9dlfYYLGQJeEx4+DwmSQLwDVXJg=" template_dir = "mine" json = true @@ -612,6 +619,7 @@ mod test { .port(7810) .workers(21) .log_level(LoggingLevel::Critical) + .keep_alive(None) .secret_key("8Xui8SN4mI+7egV/9dlfYYLGQJeEx4+DwmSQLwDVXJg=") .extra("template_dir", "mine") .extra("json", true) @@ -886,6 +894,82 @@ mod test { "#.to_string(), TEST_CONFIG_FILENAME).is_err()); } + #[test] + fn test_good_keep_alives() { + // Take the lock so changing the environment doesn't cause races. + let _env_lock = ENV_LOCK.lock().unwrap(); + env::set_var(CONFIG_ENV, "stage"); + + check_config!(RocketConfig::parse(r#" + [stage] + keep_alive = 10 + "#.to_string(), TEST_CONFIG_FILENAME), { + default_config(Staging).keep_alive(10) + }); + + check_config!(RocketConfig::parse(r#" + [stage] + keep_alive = 0 + "#.to_string(), TEST_CONFIG_FILENAME), { + default_config(Staging).keep_alive(0) + }); + + check_config!(RocketConfig::parse(r#" + [stage] + keep_alive = 348 + "#.to_string(), TEST_CONFIG_FILENAME), { + default_config(Staging).keep_alive(348) + }); + + check_config!(RocketConfig::parse(r#" + [stage] + keep_alive = false + "#.to_string(), TEST_CONFIG_FILENAME), { + default_config(Staging).keep_alive(None) + }); + + check_config!(RocketConfig::parse(r#" + [stage] + keep_alive = "none" + "#.to_string(), TEST_CONFIG_FILENAME), { + default_config(Staging).keep_alive(None) + }); + + check_config!(RocketConfig::parse(r#" + [stage] + keep_alive = "None" + "#.to_string(), TEST_CONFIG_FILENAME), { + default_config(Staging).keep_alive(None) + }); + } + + #[test] + fn test_bad_keep_alives() { + // Take the lock so changing the environment doesn't cause races. + let _env_lock = ENV_LOCK.lock().unwrap(); + env::remove_var(CONFIG_ENV); + + assert!(RocketConfig::parse(r#" + [dev] + keep_alive = true + "#.to_string(), TEST_CONFIG_FILENAME).is_err()); + + assert!(RocketConfig::parse(r#" + [dev] + keep_alive = -10 + "#.to_string(), TEST_CONFIG_FILENAME).is_err()); + + assert!(RocketConfig::parse(r#" + [dev] + keep_alive = "Some(10)" + "#.to_string(), TEST_CONFIG_FILENAME).is_err()); + + assert!(RocketConfig::parse(r#" + [dev] + keep_alive = 4294967296 + "#.to_string(), TEST_CONFIG_FILENAME).is_err()); + } + #[test] fn test_good_log_levels() { // Take the lock so changing the environment doesn't cause races. diff --git a/lib/src/config/toml_ext.rs b/lib/src/config/toml_ext.rs index ed60b2e4..103ce49d 100644 --- a/lib/src/config/toml_ext.rs +++ b/lib/src/config/toml_ext.rs @@ -4,6 +4,7 @@ use std::collections::BTreeMap; use config::Value; pub fn parse_simple_toml_value(string: &str) -> Result { + let string = string.trim(); if string.is_empty() { return Err("value is empty") } diff --git a/lib/src/rocket.rs b/lib/src/rocket.rs index 249d411a..ccd056fb 100644 --- a/lib/src/rocket.rs +++ b/lib/src/rocket.rs @@ -2,12 +2,15 @@ use std::collections::HashMap; use std::str::from_utf8; use std::cmp::min; use std::io::{self, Write}; +use std::time::Duration; use std::mem; use yansi::Paint; use state::Container; -#[cfg(feature = "tls")] use hyper_sync_rustls::TlsServer; +#[cfg(feature = "tls")] +use hyper_sync_rustls::TlsServer; + use {logger, handler}; use ext::ReadExt; use config::{self, Config, LoggedValue}; @@ -383,6 +386,11 @@ impl Rocket { launch_info_!("secret key: {}", Paint::white(&config.secret_key)); launch_info_!("limits: {}", Paint::white(&config.limits)); + match config.keep_alive { + Some(v) => launch_info_!("keep-alive: {}", Paint::white(format!("{}s", v))), + None => launch_info_!("keep-alive: {}", Paint::white("disabled")), + } + let tls_configured = config.tls.is_some(); if tls_configured && cfg!(feature = "tls") { launch_info_!("tls: {}", Paint::white("enabled")); @@ -669,6 +677,10 @@ impl Rocket { Err(e) => return LaunchError::from(e), } + // Set the keep-alive. + let timeout = self.config.keep_alive.map(|s| Duration::from_secs(s as u64)); + server.keep_alive(timeout); + // Run the launch fairings. self.fairings.handle_launch(&self);