add ETag & conditional get for rustdoc HTML pages

Author: syphar

Cargo.lock

@@ -86,6 +86,7 @@ uuid = { version = "1.1.2", features = ["v4"]}
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 serde_with = "3.4.0"
+bincode = "2.0.1"
 
 # axum dependencies
 async-trait = "0.1.83"

Cargo.toml

@@ -40,6 +40,20 @@ mod version_impl {
         }
     }
 
+    impl bincode::Encode for Version {
+        fn encode<E: bincode::enc::Encoder>(
+            &self,
+            encoder: &mut E,
+        ) -> Result<(), bincode::error::EncodeError> {
+            self.0.major.encode(encoder)?;
+            self.0.minor.encode(encoder)?;
+            self.0.patch.encode(encoder)?;
+            bincode::Encode::encode(self.0.pre.as_str(), encoder)?;
+            bincode::Encode::encode(self.0.build.as_str(), encoder)?;
+            Ok(())
+        }
+    }
+
     impl Type<Postgres> for Version {
         fn type_info() -> PgTypeInfo {
             <String as Type<Postgres>>::type_info()

src/db/types/version.rs

@@ -44,7 +44,17 @@ pub struct CrateOwner {
 }
 
 #[derive(
-    Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, sqlx::Type,
+    Debug,
+    Clone,
+    Copy,
+    PartialEq,
+    Eq,
+    PartialOrd,
+    Ord,
+    Serialize,
+    Deserialize,
+    sqlx::Type,
+    bincode::Encode,
 )]
 #[sqlx(type_name = "owner_kind", rename_all = "lowercase")]
 #[serde(rename_all = "lowercase")]

src/registry_api.rs

@@ -22,18 +22,21 @@ use crate::{
         page::TemplateData,
     },
 };
-use anyhow::Context as _;
+use anyhow::{Context as _, anyhow};
 use axum::body::Bytes;
 use axum::{Router, body::Body, http::Request, response::Response as AxumResponse};
 use axum_extra::headers::{ETag, HeaderMapExt as _};
 use fn_error_context::context;
 use futures_util::stream::TryStreamExt;
-use http::{HeaderMap, StatusCode, header::CACHE_CONTROL};
+use http::{
+    HeaderMap, HeaderName, HeaderValue, StatusCode,
+    header::{CACHE_CONTROL, CONTENT_TYPE},
+};
 use http_body_util::BodyExt;
 use opentelemetry_sdk::metrics::InMemoryMetricExporter;
 use serde::de::DeserializeOwned;
 use sqlx::Connection as _;
-use std::{fs, future::Future, panic, rc::Rc, str::FromStr, sync::Arc};
+use std::{collections::HashMap, fs, future::Future, panic, rc::Rc, str::FromStr, sync::Arc};
 use tokio::{runtime, task::block_in_place};
 use tower::ServiceExt;
 use tracing::error;
@@ -137,11 +140,13 @@ pub(crate) trait AxumRouterTestExt {
         config: &Config,
     ) -> Result<AxumResponse>;
     async fn assert_not_found(&self, path: &str) -> Result<()>;
-    async fn assert_success_and_conditional_get(
+    async fn assert_conditional_get(
         &self,
-        path: &str,
-        expected_body: &str,
+        initial_path: &str,
+        uncached_response: &AxumResponse,
     ) -> Result<()>;
+    async fn assert_success_and_conditional_get(&self, path: &str) -> Result<()>;
+
     async fn assert_success_cached(
         &self,
         path: &str,
@@ -187,37 +192,66 @@ impl AxumRouterTestExt for axum::Router {
         Ok(response)
     }
 
-    async fn assert_success_and_conditional_get(
+    async fn assert_conditional_get(
         &self,
-        path: &str,
-        expected_body: &str,
+        initial_path: &str,
+        uncached_response: &AxumResponse,
     ) -> Result<()> {
-        let etag: ETag = {
-            // uncached response
-            let response = self.assert_success(path).await?;
-            let etag: ETag = response.headers().typed_get().unwrap();
+        let etag: ETag = uncached_response
+            .headers()
+            .typed_get()
+            .ok_or_else(|| anyhow!("missing ETag header"))?;
 
-            assert_eq!(response.text().await?, expected_body);
+        let if_none_match = IfNoneMatch::from(etag.clone());
 
-            etag
-        };
+        // general rule:
+        //
+        // if a header influences how any client or intermediate proxy should treat the response,
+        // it should be repeated on the 304 response.
+        //
+        // This logic assumes _all_ headers have to be repeated, except for a few known ones.
+        const NON_CACHE_HEADERS: &[&HeaderName] = &[&CONTENT_TYPE];
 
-        let if_none_match = IfNoneMatch::from(etag.clone());
+        // store original headers, to assert that they are repeated on the 304 response.
+        let original_headers: HashMap<HeaderName, HeaderValue> = uncached_response
+            .headers()
+            .iter()
+            .filter(|(k, _)| !NON_CACHE_HEADERS.contains(k))
+            .map(|(k, v)| (k.clone(), v.clone()))
+            .collect();
 
         {
-            // cached response
-            let response = self
-                .get_with_headers(path, |headers| {
+            let cached_response = self
+                .get_with_headers(initial_path, |headers| {
+                    headers.insert(X_RLNG_SOURCE_CDN, HeaderValue::from_static("fastly"));
                     headers.typed_insert(if_none_match);
                 })
                 .await?;
-            assert_eq!(response.status(), StatusCode::NOT_MODIFIED);
-            // etag is repeated
-            assert_eq!(response.headers().typed_get::<ETag>().unwrap(), etag);
+            assert_eq!(cached_response.status(), StatusCode::NOT_MODIFIED);
+
+            // most headers are repeated on the 304 response.
+            let cached_response_headers: HashMap<HeaderName, HeaderValue> = cached_response
+                .headers()
+                .iter()
+                .filter_map(|(k, v)| {
+                    if original_headers.contains_key(k) {
+                        Some((k.clone(), v.clone()))
+                    } else {
+                        None
+                    }
+                })
+                .collect();
+
+            assert_eq!(original_headers, cached_response_headers);
         }
         Ok(())
     }
 
+    async fn assert_success_and_conditional_get(&self, path: &str) -> Result<()> {
+        self.assert_conditional_get(path, &self.assert_success(path).await?)
+            .await
+    }
+
     async fn assert_not_found(&self, path: &str) -> Result<()> {
         let response = self.get(path).await?;
 

src/test/mod.rs

@@ -135,6 +135,22 @@ pub(crate) struct Dependency {
     pub(crate) optional: bool,
 }
 
+impl bincode::Encode for Dependency {
+    fn encode<E: bincode::enc::Encoder>(
+        &self,
+        encoder: &mut E,
+    ) -> Result<(), bincode::error::EncodeError> {
+        self.name.encode(encoder)?;
+        // FIXME: VersionReq does not implement Encode, so we serialize it to string
+        // Could be fixable by wrapping VersionReq in a newtype
+        self.req.to_string().encode(encoder)?;
+        self.kind.encode(encoder)?;
+        self.rename.encode(encoder)?;
+        self.optional.encode(encoder)?;
+        Ok(())
+    }
+}
+
 impl Dependency {
     #[cfg(test)]
     pub fn new(name: String, req: VersionReq) -> Dependency {

src/utils/cargo_metadata.rs

@@ -13,7 +13,11 @@ use axum::{
     response::Response as AxumResponse,
 };
 use axum_extra::headers::HeaderMapExt as _;
-use http::{HeaderMap, HeaderName, HeaderValue, header::CACHE_CONTROL, request::Parts};
+use http::{
+    HeaderMap, HeaderName, HeaderValue, StatusCode,
+    header::{CACHE_CONTROL, ETAG},
+    request::Parts,
+};
 use serde::Deserialize;
 use std::{convert::Infallible, sync::Arc};
 use tracing::error;
@@ -235,6 +239,15 @@ pub(crate) async fn cache_middleware(
         response.headers(),
     );
 
+    debug_assert!(
+        response.status() == StatusCode::NOT_MODIFIED
+            || response.status().is_success()
+            || !response.headers().contains_key(ETAG),
+        "only successful or not-modified responses should have etags. \n{:?}\n{:?}",
+        response.status(),
+        response.headers(),
+    );
+
     // extract cache policy, default to "forbid caching everywhere".
     // We only use cache policies in our successful responses (with content, or redirect),
     // so any errors (4xx, 5xx) should always get "NoCaching".

src/web/cache.rs

@@ -19,6 +19,23 @@ pub struct EscapedURI {
     fragment: Option<String>,
 }
 
+impl bincode::Encode for EscapedURI {
+    fn encode<E: bincode::enc::Encoder>(
+        &self,
+        encoder: &mut E,
+    ) -> Result<(), bincode::error::EncodeError> {
+        // encode as separate parts so we don't have to clone
+        self.uri.scheme_str().encode(encoder)?;
+        self.uri.authority().map(|a| a.as_str()).encode(encoder)?;
+        self.uri
+            .path_and_query()
+            .map(|pq| pq.as_str())
+            .encode(encoder)?;
+        self.fragment.encode(encoder)?;
+        Ok(())
+    }
+}
+
 impl EscapedURI {
     pub fn from_uri(uri: Uri) -> Self {
         if uri.path_and_query().is_some() {

src/web/escaped_uri.rs

@@ -26,7 +26,7 @@ pub(crate) const ROOT_RUSTDOC_HTML_FILES: &[&str] = &[
     "scrape-examples-help.html",
 ];
 
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, bincode::Encode)]
 pub(crate) enum PageKind {
     Rustdoc,
     Source,
@@ -42,7 +42,7 @@ pub(crate) enum PageKind {
 /// All of these have more or less detail depending on how much metadata we have here.
 /// Maintains some additional fields containing "fixed" things, whos quality
 /// gets better the more metadata we provide.
-#[derive(Clone, PartialEq)]
+#[derive(Clone, PartialEq, bincode::Encode)]
 pub(crate) struct RustdocParams {
     // optional behaviour marker
     page_kind: Option<PageKind>,

src/web/extractors/rustdoc.rs

@@ -3,19 +3,56 @@ mod if_none_match;
 mod surrogate_key;
 
 use axum_extra::headers::ETag;
-pub use canonical_url::CanonicalUrl;
 use http::HeaderName;
+use std::io::{self, Write};
+
+pub use canonical_url::CanonicalUrl;
 pub(crate) use if_none_match::IfNoneMatch;
 pub use surrogate_key::{SURROGATE_KEY, SurrogateKey, SurrogateKeys};
 
 /// Fastly's Surrogate-Control header
 /// https://www.fastly.com/documentation/reference/http/http-headers/Surrogate-Control/
 pub static SURROGATE_CONTROL: HeaderName = HeaderName::from_static("surrogate-control");
 
+/// X-Robots-Tag header for search engines.
+pub static X_ROBOTS_TAG: HeaderName = HeaderName::from_static("x-robots-tag");
+
 /// compute our etag header value from some content
 ///
 /// Has to match the implementation in our build-script.
 pub fn compute_etag<T: AsRef<[u8]>>(content: T) -> ETag {
-    let digest = md5::compute(&content);
-    format!("\"{:x}\"", digest).parse().unwrap()
+    let mut computer = ETagComputer::new();
+    computer.write_all(content.as_ref()).unwrap();
+    computer.finalize()
+}
+
+/// Helper type to compute ETag values.
+///
+/// Works the same way as the inner `md5::Context`,
+/// but produces an `ETag` when finalized.
+pub(crate) struct ETagComputer(md5::Context);
+
+impl ETagComputer {
+    pub fn new() -> Self {
+        Self(md5::Context::new())
+    }
+
+    pub fn consume<T: AsRef<[u8]>>(&mut self, data: T) {
+        self.0.consume(data.as_ref());
+    }
+
+    pub fn finalize(self) -> ETag {
+        let digest = self.0.finalize();
+        format!("\"{:x}\"", digest).parse().unwrap()
+    }
+}
+
+impl io::Write for ETagComputer {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        self.0.write(buf)
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        self.0.flush()
+    }
 }