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@Edge
Use this solution in your Compute@Edge service:
- 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 activeconst 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.";
/// Decisionsenum 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 decisionfn 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 responsefn 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 requestfn 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 cookiefn 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 keyfn 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.