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:

  1. Rust
/// The name of a backend server associated with this service.
const BACKEND_NAME: &str = "httpbin.org";
/// 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 "auth-user-id" header or "client ip" as user identifier
let authed_user_id = match req.get_header("auth-user-id") {
Some(h) => h.to_str()?.to_string(),
None => 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
// Origin host header is configured in "Override host" of the backend configuration
// Do below if "Override host" is not configured for backend "httpbin.org"
// req.headers_mut().insert(header::HOST, HeaderValue::from_static("httpbin.org"));
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))
}