Follow redirects at the edge

[dependencies]
fastly = "0.8.0"
lazy_static = "1.4"
log = "0.4"
log-fastly = "0.8"
regex = "1.3"

5 lines
/// The name of a backend server associated with this service.
const BACKEND_NAME: &str = "my-backend";

/// Max redirect that is allow for a request
const MAX_REDIRECT_COUNT: u32 = 2;

lazy_static! {
    static ref LOCATION_CAP: Regex = Regex::new(r"^(?:https?://([^/]+))?(/.*)?$").unwrap();
}

#[fastly::main]
fn main(mut req: Request) -> Result<Response, Error> {
    log_fastly::init_simple("my_log", log::LevelFilter::Info);
    fastly::log::set_panic_endpoint("my_log").unwrap();

    // Stack the client request host header for redirect
    let request_host = req
        .get_header(header::HOST)
        .expect("host will be present")
        .clone();

    // Make a clone of client request header for redirect
    let mut req_redirect = req.clone_with_body();

    // Send request to backend
    let mut resp = req.send(BACKEND_NAME)?;

    // Do redirect if necessary
    let mut redirect_count = 0;
    while resp.get_status().is_redirection() && redirect_count < MAX_REDIRECT_COUNT {
        if let Some((redirect_host, redirect_path_query)) = get_location_from_resp(&resp) {
            log::info!(
                "Redirect to {}{}, redirect count {}",
                redirect_host,
                redirect_path_query,
                redirect_count
            );

            // Only do so if the location header does not specify a host,
            // or the host matches the client-side host header
            if redirect_host.is_empty() || redirect_host == request_host {
                // Set the path and query of the redirect request
                let redirect_url = format!("http://dummy-host{}", redirect_path_query);
                req_redirect.set_url(redirect_url);

                // Make a clone of current request in case we got another redirect from the response
                let req_clone = req_redirect.clone_with_body();

                // Send the redirect request to backend
                resp = req_redirect.send(BACKEND_NAME)?;

                // Stack info for next round of redirect
                req_redirect = req_clone;

                redirect_count += 1;

                continue;
            }
        }

        // No redirect happens
        break;
    }

    Ok(resp)
}

/// Get location information from response header
fn get_location_from_resp(resp: &Response) -> Option<(&str, &str)> {
    let location = resp.get_header_str(header::LOCATION)?;
    let location_cap = LOCATION_CAP.captures(location)?;

    let host = match location_cap.get(1) {
        Some(host_match) => host_match.as_str(),
        // The location may contain no host, but it is still a good location
        // example: "/new_location_path/abc"
        None => "",
    };

    let path_query = match location_cap.get(2) {
        Some(path_query_match) => path_query_match.as_str(),
        // The location may contain no path, but it is still a good location
        // example: "http://www.fastly.com"
        None => "/",
    };

    Some((host, path_query))
}