Flatten the curve of major traffic spikes with a waiting room

[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"

13 lines
/// 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))
}