From 6db07edcd5ceea434acfc48e005df9fa5b707ee4 Mon Sep 17 00:00:00 2001 From: Oliver Stenbom Date: Tue, 23 May 2023 08:29:18 +0200 Subject: [PATCH] Cache cache_routes using worker cache api --- Cargo.lock | 13 ++++--- linkup-cli/src/local_server.rs | 2 +- linkup/src/lib.rs | 21 +---------- linkup/src/memory_session_store.rs | 2 +- worker/Cargo.toml | 1 + worker/src/lib.rs | 60 ++++++++++++++++++++++++++---- 6 files changed, 64 insertions(+), 35 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index efeb49f..843bb15 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -214,9 +214,9 @@ dependencies = [ [[package]] name = "aho-corasick" -version = "0.7.20" +version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cc936419f96fa211c1b9166887b38e5e40b19958e5b895be7c1f93adec7071ac" +checksum = "67fc08ce920c31afb70f013dcce1bfc3a3195de6a228474e45e1f145b36f8d04" dependencies = [ "memchr", ] @@ -1066,6 +1066,7 @@ dependencies = [ "futures", "http", "linkup", + "regex", "reqwest", "worker", ] @@ -1504,9 +1505,9 @@ dependencies = [ [[package]] name = "regex" -version = "1.7.1" +version = "1.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "48aaa5748ba571fb95cd2c85c09f629215d3a6ece942baa100950af03a34f733" +checksum = "af83e617f331cc6ae2da5443c602dfa5af81e517212d9d611a5b3ba1777b5370" dependencies = [ "aho-corasick", "memchr", @@ -1515,9 +1516,9 @@ dependencies = [ [[package]] name = "regex-syntax" -version = "0.6.28" +version = "0.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "456c603be3e8d448b072f410900c09faf164fbce2d480456f50eea6e25f9c848" +checksum = "a5996294f19bd3aae0453a862ad728f60e6600695733dd5df01da90c54363a3c" [[package]] name = "reqwest" diff --git a/linkup-cli/src/local_server.rs b/linkup-cli/src/local_server.rs index 8aba9e6..663d482 100644 --- a/linkup-cli/src/local_server.rs +++ b/linkup-cli/src/local_server.rs @@ -114,7 +114,7 @@ async fn linkup_request_handler( }; let extra_resp_headers = - additional_response_headers(req.path().to_string(), config.cache_routes); + additional_response_headers(); convert_reqwest_response(response, extra_resp_headers) .await diff --git a/linkup/src/lib.rs b/linkup/src/lib.rs index 0995f70..b63d3c7 100644 --- a/linkup/src/lib.rs +++ b/linkup/src/lib.rs @@ -1,6 +1,5 @@ use async_trait::async_trait; use rand::Rng; -use regex::Regex; use std::collections::HashMap; use thiserror::Error; @@ -88,10 +87,7 @@ pub fn get_additional_headers( additional_headers } -pub fn additional_response_headers( - path: String, - cache_routes: Option>, -) -> HashMap { +pub fn additional_response_headers() -> HashMap { let mut headers = HashMap::new(); headers.insert( @@ -102,21 +98,6 @@ pub fn additional_response_headers( headers.insert("Access-Control-Allow-Headers".to_string(), "*".to_string()); headers.insert("Access-Control-Max-Age".to_string(), "86400".to_string()); - // only insert the cache-control header if the path does not match any of the cache routes - if let Some(routes) = cache_routes { - if !routes.iter().any(|route| route.is_match(&path)) { - headers.insert( - "Cache-Control".to_string(), - "no-store".to_string(), - ); - } - } else { - headers.insert( - "Cache-Control".to_string(), - "no-store".to_string(), - ); - } - headers } diff --git a/linkup/src/memory_session_store.rs b/linkup/src/memory_session_store.rs index e67a247..8a9a10e 100644 --- a/linkup/src/memory_session_store.rs +++ b/linkup/src/memory_session_store.rs @@ -1,6 +1,6 @@ use std::{ collections::HashMap, - sync::{Mutex, RwLock}, + sync::{RwLock}, }; use async_trait::async_trait; diff --git a/worker/Cargo.toml b/worker/Cargo.toml index 53d43d2..2a33753 100644 --- a/worker/Cargo.toml +++ b/worker/Cargo.toml @@ -22,6 +22,7 @@ futures = "0.3" console_error_panic_hook = { version = "0.1.1", optional = true } http = "0.2.9" reqwest = "0.11.17" +regex = "1.8.1" # [profile.release] # Tell `rustc` to optimize for small code size. diff --git a/worker/src/lib.rs b/worker/src/lib.rs index 9e78b6d..4a97f96 100644 --- a/worker/src/lib.rs +++ b/worker/src/lib.rs @@ -1,3 +1,4 @@ +use regex::Regex; use std::{collections::HashMap, sync::Arc}; use kv_store::CfWorkerStringStore; @@ -54,12 +55,45 @@ async fn linkup_session_handler(mut req: Request, sessions: SessionAllocator) -> } } -async fn linkup_request_handler(mut req: Request, sessions: SessionAllocator) -> Result { - let body_bytes = match req.bytes().await { - Ok(bytes) => bytes, - Err(_) => return plaintext_error("Bad or missing request body", 400), - }; +async fn get_cached_req( + req: &Request, + cache_routes: &Option>, +) -> Result> { + let path = req.path(); + + if let Some(routes) = cache_routes { + if routes.iter().any(|route| route.is_match(&path)) { + let url = req.url()?; + Cache::default().get(url.to_string(), false).await + } else { + Ok(None) + } + } else { + Ok(None) + } +} + +async fn set_cached_req( + req: &Request, + mut resp: Response, + cache_routes: Option>, +) -> Result { + let path = req.path(); + + if let Some(routes) = cache_routes { + if routes.iter().any(|route| route.is_match(&path)) { + let url = req.url()?; + let cache_resp = resp.cloned()?; + Cache::default().put(url.to_string(), cache_resp).await?; + + return Ok(resp); + } + } + Ok(resp) +} + +async fn linkup_request_handler(mut req: Request, sessions: SessionAllocator) -> Result { let url = match req.url() { Ok(url) => url.to_string(), Err(_) => return plaintext_error("Bad or missing request url", 400), @@ -77,6 +111,15 @@ async fn linkup_request_handler(mut req: Request, sessions: SessionAllocator) -> Err(_) => return plaintext_error("Could not find a linkup session for this request. Use a linkup subdomain or context headers like Referer/tracestate", 422), }; + if let Some(cached_response) = get_cached_req(&req, &config.cache_routes).await? { + return Ok(cached_response); + } + + let body_bytes = match req.bytes().await { + Ok(bytes) => bytes, + Err(_) => return plaintext_error("Bad or missing request body", 400), + }; + let destination_url = match get_target_url(url.clone(), headers.clone(), &config, &session_name) { Some(result) => result, @@ -104,9 +147,12 @@ async fn linkup_request_handler(mut req: Request, sessions: SessionAllocator) -> Err(e) => return plaintext_error(format!("Failed to proxy request: {}", e), 502), }; - let extra_resp_headers = additional_response_headers(req.path(), config.cache_routes); - convert_reqwest_response_to_cf(response, extra_resp_headers).await + let mut cf_resp = convert_reqwest_response_to_cf(response, additional_response_headers()).await?; + + cf_resp = set_cached_req(&req, cf_resp, config.cache_routes).await?; + + Ok(cf_resp) } async fn linkup_ws_handler(req: Request, sessions: SessionAllocator) -> Result {