Merge pull request #21 from dfinity/igor/fix-acme

Quick fix re-export

Author: blind-oracle

.github/workflows/test.yml

@@ -18,5 +18,17 @@ jobs:
 
       - uses: Swatinem/rust-cache@9d47c6ad4b02e050fd481d890b2ea34778fd09d6
 
-      - name: Run tests
-        run: cargo test
+      - name: Run tests (all features)
+        run: cargo test --all-features
+
+      - name: Run tests (acme_dns)
+        run: cargo test --features acme_dns
+
+      - name: Run tests (acme_alpn)
+        run: cargo test --features acme_alpn
+
+      - name: Run tests (vector)
+        run: cargo test --features vector
+
+      - name: Run tests (clients-hyper)
+        run: cargo test --features clients-hyper

Cargo.toml

@@ -103,7 +103,9 @@ tower = { version = "0.5.1", features = ["util"] }
 tower-service = "0.3.3"
 tracing = "0.1.40"
 url = "2.5.3"
-uuid = { version = "1.16.0", features = ["v7"] }
+# DO NOT upgrade, this breaks monorepo compatibility
+# Read https://github.com/uuid-rs/uuid/releases/tag/1.13.0
+uuid = { version = "=1.12.1", features = ["v7"] }
 vrl = { version = "0.23.0", default-features = false, features = [
     "value",
 ], optional = true }

src/http/cache.rs

@@ -31,33 +31,52 @@ use strum_macros::{Display, IntoStaticStr};
 use tokio::{select, sync::Mutex, time::sleep};
 use tokio_util::sync::CancellationToken;
 
+use super::{Error as HttpError, body::buffer_body, calc_headers_size, extract_authority};
 use crate::{http::headers::X_CACHE_TTL, tasks::Run};
 
-use super::{Error as HttpError, body::buffer_body, calc_headers_size, extract_authority};
+pub trait CustomBypassReason:
+    Debug + Clone + std::fmt::Display + Into<&'static str> + PartialEq + Eq + Send + Sync + 'static
+{
+}
+
+#[derive(Debug, Clone, Display, PartialEq, Eq, IntoStaticStr)]
+pub enum CustomBypassReasonDummy {}
+impl CustomBypassReason for CustomBypassReasonDummy {}
 
 #[derive(Debug, Clone, Display, PartialEq, Eq, IntoStaticStr)]
 #[strum(serialize_all = "snake_case")]
-pub enum CacheBypassReason {
+pub enum CacheBypassReason<R: CustomBypassReason> {
     MethodNotCacheable,
     SizeUnknown,
     BodyTooBig,
     HTTPError,
     UnableToExtractKey,
+    UnableToRunBypasser,
     CacheControl,
+    Custom(R),
+}
+
+impl<R: CustomBypassReason> CacheBypassReason<R> {
+    pub fn into_str(self) -> &'static str {
+        match self {
+            Self::Custom(v) => v.into(),
+            _ => self.into(),
+        }
+    }
 }
 
 #[derive(Debug, Clone, Display, PartialEq, Eq, Default, IntoStaticStr)]
 #[strum(serialize_all = "SCREAMING_SNAKE_CASE")]
-pub enum CacheStatus {
+pub enum CacheStatus<R: CustomBypassReason = CustomBypassReasonDummy> {
     #[default]
     Disabled,
-    Bypass(CacheBypassReason),
+    Bypass(CacheBypassReason<R>),
     Hit,
     Miss,
 }
 
 // Injects itself into a given response to be accessible by middleware
-impl CacheStatus {
+impl<B: CustomBypassReason> CacheStatus<B> {
     pub fn with_response<T>(self, mut resp: Response<T>) -> Response<T> {
         resp.extensions_mut().insert(self);
         resp
@@ -68,6 +87,8 @@ impl CacheStatus {
 pub enum Error {
     #[error("unable to extract key from request: {0}")]
     ExtractKey(String),
+    #[error("unable to execute bypasser: {0}")]
+    ExecuteBypasser(String),
     #[error("timed out while fetching body")]
     FetchBodyTimeout,
     #[error("body is too big")]
@@ -80,9 +101,9 @@ pub enum Error {
     Other(String),
 }
 
-enum ResponseType {
+enum ResponseType<R: CustomBypassReason> {
     Fetched(Response<Bytes>, Duration),
-    Streamed(Response, CacheBypassReason),
+    Streamed(Response, CacheBypassReason<R>),
 }
 
 #[derive(Clone)]
@@ -121,6 +142,26 @@ pub trait KeyExtractor: Clone + Send + Sync + Debug + 'static {
     fn extract<T>(&self, req: &Request<T>) -> Result<Self::Key, Error>;
 }
 
+/// Trait to decide if we need to bypass caching of the given request
+pub trait Bypasser: Clone + Send + Sync + Debug + 'static {
+    /// Custom bypass reason
+    type BypassReason: CustomBypassReason;
+
+    /// Checks if we should bypass the given request
+    fn bypass<T>(&self, req: &Request<T>) -> Result<Option<Self::BypassReason>, Error>;
+}
+
+#[derive(Debug, Clone)]
+pub struct NoopBypasser;
+
+impl Bypasser for NoopBypasser {
+    type BypassReason = CustomBypassReasonDummy;
+
+    fn bypass<T>(&self, _req: &Request<T>) -> Result<Option<Self::BypassReason>, Error> {
+        Ok(None)
+    }
+}
+
 #[derive(Clone)]
 pub struct Metrics {
     lock_await: HistogramVec,
@@ -245,8 +286,12 @@ fn infer_ttl<T>(req: &Response<T>) -> Option<CacheControl> {
         if ["no-cache", "no-store"].contains(&k) {
             Some(CacheControl::NoCache)
         } else if k == "max-age" {
-            v.and_then(|x| x.parse::<u64>().ok())
-                .map(|x| CacheControl::MaxAge(Duration::from_secs(x)))
+            let v = v.and_then(|x| x.parse::<u64>().ok());
+            if v == Some(0) {
+                Some(CacheControl::NoCache)
+            } else {
+                v.map(|x| CacheControl::MaxAge(Duration::from_secs(x)))
+            }
         } else {
             None
         }
@@ -268,16 +313,29 @@ impl<K: KeyExtractor> Expiry<K::Key, Arc<Entry>> for Expirer<K> {
 }
 
 /// Builds a cache using some overridable defaults
-pub struct CacheBuilder<K: KeyExtractor> {
+pub struct CacheBuilder<K: KeyExtractor, B: Bypasser> {
     key_extractor: K,
+    bypasser: Option<B>,
     opts: Opts,
     registry: Registry,
 }
 
-impl<K: KeyExtractor> CacheBuilder<K> {
+impl<K: KeyExtractor> CacheBuilder<K, NoopBypasser> {
     pub fn new(key_extractor: K) -> Self {
         Self {
             key_extractor,
+            bypasser: None,
+            opts: Opts::default(),
+            registry: Registry::new(),
+        }
+    }
+}
+
+impl<K: KeyExtractor, B: Bypasser> CacheBuilder<K, B> {
+    pub fn new_with_bypasser(key_extractor: K, bypasser: B) -> Self {
+        Self {
+            key_extractor,
+            bypasser: Some(bypasser),
             opts: Opts::default(),
             registry: Registry::new(),
         }
@@ -344,15 +402,16 @@ impl<K: KeyExtractor> CacheBuilder<K> {
     }
 
     /// Try to build the cache from this builder
-    pub fn build(self) -> Result<Cache<K>, Error> {
-        Cache::new(self.opts, self.key_extractor, &self.registry)
+    pub fn build(self) -> Result<Cache<K, B>, Error> {
+        Cache::new(self.opts, self.key_extractor, self.bypasser, &self.registry)
     }
 }
 
-pub struct Cache<K: KeyExtractor> {
+pub struct Cache<K: KeyExtractor, B: Bypasser = NoopBypasser> {
     store: MokaCache<K::Key, Arc<Entry>, RandomState>,
     locks: MokaCache<K::Key, Arc<Mutex<()>>, RandomState>,
     key_extractor: K,
+    bypasser: Option<B>,
     metrics: Metrics,
     opts: Opts,
 }
@@ -366,8 +425,13 @@ fn weigh_entry<K: KeyExtractor>(_k: &K::Key, v: &Arc<Entry>) -> u32 {
     size as u32
 }
 
-impl<K: KeyExtractor + 'static> Cache<K> {
-    pub fn new(opts: Opts, key_extractor: K, registry: &Registry) -> Result<Self, Error> {
+impl<K: KeyExtractor + 'static, B: Bypasser + 'static> Cache<K, B> {
+    pub fn new(
+        opts: Opts,
+        key_extractor: K,
+        bypasser: Option<B>,
+        registry: &Registry,
+    ) -> Result<Self, Error> {
         if opts.max_item_size as u64 >= opts.cache_size {
             return Err(Error::Other(
                 "Cache item size should be less than whole cache size".into(),
@@ -390,6 +454,7 @@ impl<K: KeyExtractor + 'static> Cache<K> {
                 .build_with_hasher(RandomState::default()),
 
             key_extractor,
+            bypasser,
             metrics: Metrics::new(registry),
 
             opts,
@@ -434,11 +499,11 @@ impl<K: KeyExtractor + 'static> Cache<K> {
         let (cache_status, response) = self.process_inner(now, request, next).await?;
 
         // Record metrics
-        let cache_bypass_reason_str: &'static str = match &cache_status {
-            CacheStatus::Bypass(v) => v.into(),
+        let cache_status_str: &'static str = (&cache_status).into();
+        let cache_bypass_reason_str: &'static str = match cache_status.clone() {
+            CacheStatus::Bypass(v) => v.into_str(),
             _ => "none",
         };
-        let cache_status_str: &'static str = (&cache_status).into();
 
         let labels = &[cache_status_str, cache_bypass_reason_str];
 
@@ -456,7 +521,26 @@ impl<K: KeyExtractor + 'static> Cache<K> {
         now: Instant,
         request: Request,
         next: Next,
-    ) -> Result<(CacheStatus, Response), Error> {
+    ) -> Result<(CacheStatus<B::BypassReason>, Response), Error> {
+        // Check if we have bypasser configured
+        if let Some(b) = &self.bypasser {
+            // Run it
+            if let Ok(v) = b.bypass(&request) {
+                // If it decided to bypass - return the custom reason
+                if let Some(r) = v {
+                    return Ok((
+                        CacheStatus::Bypass(CacheBypassReason::Custom(r)),
+                        next.run(request).await,
+                    ));
+                }
+            } else {
+                return Ok((
+                    CacheStatus::Bypass(CacheBypassReason::UnableToRunBypasser),
+                    next.run(request).await,
+                ));
+            }
+        }
+
         // Check the method
         if !self.opts.methods.contains(request.method()) {
             return Ok((
@@ -526,7 +610,11 @@ impl<K: KeyExtractor + 'static> Cache<K> {
     }
 
     // Passes the request down the line and conditionally fetches the response body
-    async fn pass_request(&self, request: Request, next: Next) -> Result<ResponseType, Error> {
+    async fn pass_request(
+        &self,
+        request: Request,
+        next: Next,
+    ) -> Result<ResponseType<B::BypassReason>, Error> {
         // Execute the response & get the headers
         let response = next.run(request).await;
 
@@ -621,7 +709,7 @@ impl<K: KeyExtractor + 'static> Cache<K> {
 }
 
 #[async_trait]
-impl<K: KeyExtractor> Run for Cache<K> {
+impl<K: KeyExtractor, B: Bypasser> Run for Cache<K, B> {
     async fn run(&self, _: CancellationToken) -> Result<(), anyhow::Error> {
         self.store.run_pending_tasks();
         self.metrics.memory.set(self.store.weighted_size() as i64);
@@ -748,6 +836,25 @@ mod tests {
             .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR.into_response())
     }
 
+    #[test]
+    fn test_bypass_reason_serialize() {
+        #[derive(Debug, Clone, Display, PartialEq, Eq, IntoStaticStr)]
+        #[strum(serialize_all = "snake_case")]
+        enum CustomReasonTest {
+            Bar,
+        }
+        impl CustomBypassReason for CustomReasonTest {}
+
+        let a: CacheBypassReason<CustomReasonTest> =
+            CacheBypassReason::Custom(CustomReasonTest::Bar);
+        let txt = a.into_str();
+        assert_eq!(txt, "bar");
+
+        let a: CacheBypassReason<CustomReasonTest> = CacheBypassReason::BodyTooBig;
+        let txt = a.into_str();
+        assert_eq!(txt, "body_too_big");
+    }
+
     #[test]
     fn test_key_extractor_uri_range() {
         let x = KeyExtractorUriRange;
@@ -823,6 +930,8 @@ mod tests {
             infer_ttl(&req),
             Some(CacheControl::MaxAge(Duration::from_secs(86400)))
         );
+        req.headers_mut().insert(CACHE_CONTROL, hval!("max-age=0"));
+        assert_eq!(infer_ttl(&req), Some(CacheControl::NoCache));
 
         req.headers_mut()
             .insert(CACHE_CONTROL, hval!("max-age=foo"));

src/http/client/mod.rs

@@ -50,10 +50,9 @@ pub trait ClientWithStats: Client + Stats {
     fn to_client(self: Arc<Self>) -> Arc<dyn Client>;
 }
 
-pub trait CloneableDnsResolver:
-    Resolve + Service<HyperName> + Clone + fmt::Debug + 'static
-{
-}
+pub trait CloneableDnsResolver: Resolve + Clone + fmt::Debug + 'static {}
+
+pub trait CloneableHyperDnsResolver: Service<HyperName> + Clone + fmt::Debug + 'static {}
 
 #[derive(Clone, Debug)]
 struct Metrics {

src/http/dns.rs

@@ -21,7 +21,10 @@ use reqwest::dns::{Addrs, Name, Resolve, Resolving};
 use strum_macros::EnumString;
 use tower::Service;
 
-use super::{Error, client::CloneableDnsResolver};
+use super::{
+    Error,
+    client::{CloneableDnsResolver, CloneableHyperDnsResolver},
+};
 
 #[derive(Clone, Copy, Debug, EnumString)]
 #[strum(serialize_all = "snake_case")]
@@ -58,6 +61,7 @@ impl Default for Options {
 #[derive(Debug, Clone)]
 pub struct Resolver(Arc<TokioResolver>);
 impl CloneableDnsResolver for Resolver {}
+impl CloneableHyperDnsResolver for Resolver {}
 
 impl Resolver {
     /// Creates a new resolver with given options.

src/lib.rs

@@ -1,3 +1,5 @@
+// Needed for certain macros
+#![recursion_limit = "256"]
 #![warn(clippy::nursery)]
 #![warn(tail_expr_drop_order)]
 
@@ -8,6 +10,8 @@ pub mod types;
 #[cfg(feature = "vector")]
 pub mod vector;
 
+pub use prometheus;
+
 /// Generic error
 #[derive(thiserror::Error, Debug)]
 pub enum Error {

src/tls/acme/mod.rs

@@ -12,6 +12,7 @@ use async_trait::async_trait;
 use derive_new::new;
 use strum_macros::{Display, EnumString};
 
+#[cfg(feature = "acme_dns")]
 pub use instant_acme;
 
 #[derive(Clone, Display, EnumString, PartialEq, Eq)]