Flatten the curve of major traffic spikes with a waiting room

A totally stateless solution to hold back new users for a minimum waiting period to smooth out spikes in traffic.

VCL

Use this solution in your VCL service (click RUN below to test this solution or clone it to make changes):

Compute

Use this solution in your Compute service:

  1. Rust
Cargo.toml
Rust
[dependencies]
anyhow = "1.0.28"
base64 = "0.12"
chrono = "0.4"
fastly = "0.7.0"
hex = "0.4"
lazy_static = "1.4"
log = "0.4"
log-fastly = "0.2"
rand = "0.7"
rust-crypto-wasm = "0.3"
serde_urlencoded = "0.7.0"
main.rs
Rust
/// The name of a backend server associated with this service.
const BACKEND_NAME: &str = "httpbin";
/// Whether the waiting room is active
const CONFIG_ENABLED: bool = true;
const CONFIG_ALLOW_PERCENTAGE: u32 = 50;
const CONFIG_KEY_PAIRS: [(&str, &str); 2] = [("key1", "secret"), ("key2", "another secret")];
const CONFIG_ALLOW_PERIOD_TIMEOUT: i64 = 3600;
const CONFIG_WAIT_PERIOD_TIMEOUT: i64 = 30;
const CONFIG_ACTIVE_KEY: &str = "key1";
const CONFIG_COOKIE_LIFE_TIME: u32 = 7200;
const CONFIG_MSG_WAIT: &str = "Sorry, you have to wait";
const CONFIG_MSG_KEEP_WAITING: &str = "Please continue to wait";
const CONFIG_MSG_DENY: &str = "Sorry, we're closed right now. Please try again later.";
/// Decisions
enum Decision {
/// Request is allowed to access backend
Allow,
/// Request is denied to access backend
Deny,
/// Request is anonymous (with no valid waiting room cookie)
Anon,
/// Request is put to wait state
Wait,
/// Request is put to wait state again
ReWait,
}
impl FromStr for Decision {
type Err = Error;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s {
"allow" => Ok(Decision::Allow),
"deny" => Ok(Decision::Deny),
"anon" => Ok(Decision::Anon),
"wait" => Ok(Decision::Wait),
"re-wait" => Ok(Decision::ReWait),
_ => Ok(Decision::Anon),
}
}
}
impl fmt::Display for Decision {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
let decision_str = match *self {
Decision::Allow => "allow",
Decision::Deny => "deny",
Decision::Anon => "anon",
Decision::Wait => "wait",
Decision::ReWait => "re-wait",
};
write!(f, "{}", decision_str)
}
}
lazy_static! {
pub static ref KEY_PAIRS: HashMap<&'static str, &'static str> = {
let mut pairs = HashMap::new();
for pair in &CONFIG_KEY_PAIRS {
pairs.insert(pair.0, pair.1);
}
pairs
};
}
#[fastly::main]
fn main(req: Request) -> Result<Response, Error> {
// Log initial setup
log_fastly::init_simple("my_log", log::LevelFilter::Info);
fastly::log::set_panic_endpoint("my_log")?;
if !CONFIG_ENABLED {
// No waiting room is needed
return Ok(req.send(BACKEND_NAME)?);
}
// Identify the user, use "client ip" as user identifier
let authed_user_id = req
.get_client_ip_addr()
.ok_or_else(|| anyhow!("Failed to get client ip"))?
.to_string();
// Make decision if we should allow the request
let decision = make_waiting_decision(&authed_user_id, &req)?;
log::info!(
"Waiting room state {} for user {}",
decision,
authed_user_id
);
// Generate the cookie to set on the response
let new_cookie_option = make_waiting_room_cookie(&authed_user_id, &decision);
// Response to client
match response_to_client(req, decision, new_cookie_option) {
Ok(resp) => Ok(resp),
Err(e) => {
log::error!("error sending response to client: {}", e);
Err(e)
}
}
}
/// Response to client according to the waiting decision
fn response_to_client(
req: Request,
decision: Decision,
new_cookie_option: Option<String>,
) -> Result<Response> {
match decision {
Decision::Allow => {
// Forward request to origin
// (Ensure that the origin host header is configured in "Override host" of the backend configuration)
let mut resp = req.send(BACKEND_NAME)?;
// Set the new cookie if needed
if let Some(new_cookie) = new_cookie_option {
resp.set_header(header::SET_COOKIE, new_cookie);
}
Ok(resp)
}
_ => {
// Do not allow the origin access yet
// Compose a response to client
let mut resp = Response::from_status(StatusCode::OK)
.with_header(header::CACHE_CONTROL, "no-store, private");
// Set refresh header when waiting is needed
if matches!(decision, Decision::Anon | Decision::Wait | Decision::ReWait) {
let refresh_value = if let Some(query) = req.get_query_str() {
format!("30; url={}?{}", req.get_path(), query)
} else {
format!("30; url={}", req.get_path())
};
resp.set_header(header::REFRESH, refresh_value);
}
if let Some(new_cookie) = new_cookie_option {
resp.set_header(header::SET_COOKIE, new_cookie);
}
// Choose response message according to the wait status
let resp_body = match decision {
Decision::Anon => CONFIG_MSG_WAIT,
Decision::Wait | Decision::ReWait => CONFIG_MSG_KEEP_WAITING,
_ => CONFIG_MSG_DENY,
};
Ok(resp.with_body(resp_body))
}
}
}
/// Make new cookie for the response
fn make_waiting_room_cookie(authed_user_id: &str, decision: &Decision) -> Option<String> {
match decision {
Decision::Anon | Decision::Allow | Decision::ReWait => {
let (duration, decision_cookie) = if matches!(*decision, Decision::Allow) {
(CONFIG_ALLOW_PERIOD_TIMEOUT, Decision::Allow)
} else {
(CONFIG_WAIT_PERIOD_TIMEOUT, Decision::Wait)
};
let expires = if matches!(*decision, Decision::Allow) {
Utc::now().timestamp() + duration
} else {
// For waiting users, set the expiry time to the next boundary
let timestamp = Utc::now().timestamp();
let boundary_start = timestamp - (timestamp % duration);
boundary_start + duration * 2
};
let string_to_sign = format!(
"dec={}&exp={}&uid={}&kid={}",
decision_cookie, expires, authed_user_id, CONFIG_ACTIVE_KEY
);
let key_secret = KEY_PAIRS.get(CONFIG_ACTIVE_KEY)?;
let sig = sign(key_secret.as_bytes(), &string_to_sign);
let sig_str = hex::encode(sig);
let waitingroom_info_base64 = base64::encode_config(
format!("{}&sig={}", string_to_sign, sig_str),
base64::URL_SAFE,
);
Some(format!(
"waiting_room={}; path=/; max-age={}",
waitingroom_info_base64, CONFIG_COOKIE_LIFE_TIME
))
}
Decision::Deny => {
Some("waiting_room=deleted; path=/; expires=Thu, 01 Jan 1970 00:00:00 GMT".to_string())
}
Decision::Wait => None,
}
}
/// Make a decision for the request
fn make_waiting_decision(authed_user_id: &str, req: &Request) -> Result<Decision> {
if CONFIG_ALLOW_PERCENTAGE >= 100 {
return Ok(Decision::Allow);
}
let cookie_base64 = match get_waiting_room_cookie(req) {
Ok(cookie) => cookie,
Err(_) => {
// Special case for if user does not have a waiting_room token in cookie
return Ok(Decision::Anon);
}
};
// extract data from cookie string
let cookie_decoded = base64::decode_config(cookie_base64, base64::URL_SAFE)?;
let cookie = str::from_utf8(&cookie_decoded)?;
log::info!("Waiting Info: {}", cookie);
let qs: HashMap<&str, &str> = serde_urlencoded::from_str(cookie)?;
let expire_cookie = qs.get("exp").unwrap_or(&"");
let sig_cookie = qs.get("sig").unwrap_or(&"");
let key_id_cookie = qs.get("kid").unwrap_or(&"");
let user_id_cookie = qs.get("uid").unwrap_or(&"");
let decision_cookie_str = qs.get("dec").unwrap_or(&"");
let decision_cookie: Decision = decision_cookie_str.parse()?;
if authed_user_id != *user_id_cookie {
log::info!(
"User {} denied while using a token generated for user {}",
authed_user_id,
user_id_cookie
);
// Reset the decision to anon, as if the user didn't have a cookie
return Ok(Decision::Anon);
}
let key_secret = match KEY_PAIRS.get(key_id_cookie) {
Some(secret) => secret,
None => {
log::info!(
"Unable to check signature due to missing key {}",
key_id_cookie
);
return Ok(Decision::Anon);
}
};
let string_to_sign = format!(
"dec={}&exp={}&uid={}&kid={}",
decision_cookie_str, expire_cookie, user_id_cookie, key_id_cookie
);
let sig = sign(key_secret.as_bytes(), &string_to_sign);
let expected_sig = hex::encode(sig);
if !secure_equal(sig_cookie.as_bytes(), expected_sig.as_bytes()) {
log::info!("Invalid signature");
return Ok(Decision::Anon);
}
let expire_time: i64 = expire_cookie.parse()?;
let now = chrono::Utc::now().timestamp();
if expire_time >= now {
// Cookie has not expired yet. Keep the current decision
return Ok(decision_cookie);
}
// Actions for cookies that have reached their 'expiry time'
match decision_cookie {
Decision::Allow => {
log::info!("Expired allow token reverted to anon");
Ok(Decision::Anon)
}
Decision::Wait => {
// If the user is waiting, they've now waited their turn
// so reveal the cookie decision
let mut rng = rand::thread_rng();
let idx: u32 = rng.gen_range(0, 100);
if idx < CONFIG_ALLOW_PERCENTAGE {
Ok(Decision::Allow)
} else {
Ok(Decision::ReWait)
}
}
_ => Ok(decision_cookie),
}
}
/// Get permission string from request cookie
fn get_waiting_room_cookie(req: &Request) -> Result<String> {
let cookie_val: &str = req
.get_header_str(header::COOKIE)
.ok_or_else(|| anyhow!("No cookie found"))?;
// Split at ";" not "; ", in case the cookie is ending with ";"
cookie_val
.split(';')
.find_map(|kv| {
let index = kv.find('=')?;
let (mut key, value) = kv.split_at(index);
key = key.trim();
if key != "waiting_room" {
return None;
}
// remove the "="
let value = value[1..].to_string();
Some(value)
})
.ok_or_else(|| anyhow!("No permission found in cookie"))
}
/// Generate HMAC hash of message with key
fn sign(key: &[u8], message: &str) -> Vec<u8> {
log::info!("Message: {}", message);
let mut mac = Hmac::new(Sha256::new(), key);
mac.input(message.as_bytes());
mac.result().code().to_vec()
}
/// Compare two byte slices with constant execution time. This provides some protection
/// against timing side-channel attacks.
fn secure_equal(a: &[u8], b: &[u8]) -> bool {
a.len() == b.len() && 0 == a.iter().zip(b).fold(0, |r, (a, b)| r | (a ^ b))
}
This page is part of a series in the Rate limiting topic.

User contributed notes

BETA

Do you see an error in this page? Do have an interesting use case, example or edge case people should know about? Share your knowledge and help people who are reading this page! (Comments are moderated; for support, please contact support@fastly.com)