add content-length to streamed file downloads for browsers / progress

Author: syphar

crates/bin/docs_rs_web/src/file.rs

@@ -10,7 +10,7 @@ use axum::{
 };
 use axum_extra::{
     TypedHeader,
-    headers::{ContentType, LastModified},
+    headers::{ContentLength, ContentType, LastModified},
 };
 use docs_rs_headers::IfNoneMatch;
 use docs_rs_storage::{AsyncStorage, Blob, StreamingBlob};
@@ -81,7 +81,17 @@ impl StreamingFile {
         // Future optimization could be:
         // * only forbid fastly to store, and browsers still could.
         // * implement segmented caching for large files somehow.
-        if self.0.content_length > FASTLY_CACHE_MAX_OBJECT_SIZE
+        //
+        // With this logic this error is only raised when we return the raw stream to the user,
+        // so for zip-archives & rustdoc-json files, which are also the cases where we did hit these
+        // fastly limits.
+        // Streams we decompress (rustdoc content, assets) would lead to errors in fastly if the
+        // decompressed file is >100 MiB.
+        if self
+            .0
+            .content_length
+            .map(|cl| cl > FASTLY_CACHE_MAX_OBJECT_SIZE)
+            .unwrap_or(false)
             && !matches!(cache_policy, CachePolicy::NoStoreMustRevalidate)
         {
             warn!(
@@ -113,6 +123,9 @@ impl StreamingFile {
             (
                 StatusCode::OK,
                 TypedHeader(ContentType::from(self.0.mime)),
+                self.0
+                    .content_length
+                    .map(|cl| TypedHeader(ContentLength(cl as u64))),
                 TypedHeader(last_modified),
                 self.0.etag.map(TypedHeader),
                 Extension(cache_policy),
@@ -132,7 +145,7 @@ mod tests {
     use docs_rs_headers::compute_etag;
     use docs_rs_storage::StorageKind;
     use docs_rs_types::CompressionAlgorithm;
-    use http::header::{CACHE_CONTROL, ETAG, LAST_MODIFIED};
+    use http::header::{CACHE_CONTROL, CONTENT_LENGTH, ETAG, LAST_MODIFIED};
     use std::{io, rc::Rc};
 
     const CONTENT: &[u8] = b"Hello, world!";
@@ -148,15 +161,15 @@ mod tests {
             date_updated: Utc::now(),
             compression: alg,
             etag: Some(compute_etag(&content)),
-            content_length: content.len(),
+            content_length: Some(content.len()),
             content: Box::new(io::Cursor::new(content)),
         }
     }
 
     #[test]
     fn test_big_file_stream_drops_cache_policy() {
         let mut stream = streaming_blob(CONTENT, None);
-        stream.content_length = FASTLY_CACHE_MAX_OBJECT_SIZE + 1;
+        stream.content_length = Some(FASTLY_CACHE_MAX_OBJECT_SIZE + 1);
 
         let response =
             StreamingFile(stream).into_response(None, CachePolicy::ForeverInCdnAndBrowser);
@@ -177,6 +190,14 @@ mod tests {
             let resp = stream.into_response(None, STATIC_ASSET_CACHE_POLICY);
             assert!(resp.status().is_success());
             assert!(resp.headers().get(CACHE_CONTROL).is_none());
+            assert_eq!(
+                resp.headers()
+                    .get(CONTENT_LENGTH)
+                    .unwrap()
+                    .to_str()
+                    .unwrap(),
+                CONTENT.len().to_string()
+            );
             let cache = resp
                 .extensions()
                 .get::<CachePolicy>()
@@ -209,6 +230,26 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_stream_into_response_without_content_length() -> Result<()> {
+        let mut stream = streaming_blob(CONTENT, None);
+        stream.content_length = None;
+
+        let stream_file = StreamingFile(stream);
+        let resp = stream_file.into_response(None, STATIC_ASSET_CACHE_POLICY);
+        assert!(resp.status().is_success());
+        assert!(resp.headers().get(CACHE_CONTROL).is_none());
+        assert!(resp.headers().get(CONTENT_LENGTH).is_none());
+        let cache = resp
+            .extensions()
+            .get::<CachePolicy>()
+            .expect("missing cache response extension");
+        assert!(matches!(cache, CachePolicy::ForeverInCdnAndBrowser));
+        assert!(resp.headers().get(LAST_MODIFIED).is_some());
+
+        Ok(())
+    }
+
     #[tokio::test(flavor = "multi_thread")]
     async fn file_roundtrip_axum() -> Result<()> {
         let env = TestEnvironment::new().await?;

crates/bin/docs_rs_web/src/handlers/rustdoc.rs

@@ -1104,13 +1104,26 @@ mod test {
         testing::{KRATE, V2},
     };
     use docs_rs_uri::encode_url_path;
+    use http::header::CONTENT_LENGTH;
     use kuchikiki::traits::TendrilSink;
     use pretty_assertions::assert_eq;
     use reqwest::StatusCode;
     use std::{collections::BTreeMap, str::FromStr as _};
     use test_case::test_case;
     use tracing::info;
 
+    fn has_content_len(headers: &HeaderMap) -> bool {
+        let content_length: usize = headers
+            .get(CONTENT_LENGTH)
+            .unwrap()
+            .to_str()
+            .unwrap()
+            .parse()
+            .unwrap();
+
+        content_length > 0
+    }
+
     async fn try_latest_version_redirect(
         krate: &str,
         path: &str,
@@ -3229,6 +3242,7 @@ mod test {
             resp.headers().get(CONTENT_DISPOSITION).unwrap(),
             "attachment; filename=\"rustdoc-dummy-0.2.0.zip\""
         );
+        assert!(has_content_len(resp.headers()));
         web.assert_conditional_get(path, &resp).await?;
 
         check_archive_consistency(&web.assert_success(path).await?.bytes().await?)?;
@@ -3627,6 +3641,9 @@ mod test {
                 expected_compression.file_extension()
             )
         );
+
+        assert!(has_content_len(resp.headers()));
+
         web.assert_conditional_get(&path, &resp).await?;
 
         {

crates/bin/docs_rs_web/src/testing/axum_helpers.rs

@@ -4,6 +4,7 @@ use axum::body::Bytes;
 use axum::{body::Body, http::Request, response::Response as AxumResponse};
 use axum_extra::headers::{ETag, HeaderMapExt as _};
 use docs_rs_headers::{IfNoneMatch, SURROGATE_CONTROL, SurrogateKeys};
+use http::header::CONTENT_LENGTH;
 use http::{
     HeaderMap, HeaderName, HeaderValue, StatusCode,
     header::{CACHE_CONTROL, CONTENT_TYPE},
@@ -166,7 +167,7 @@ impl AxumRouterTestExt for axum::Router {
         // 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];
+        const NON_CACHE_HEADERS: &[&HeaderName] = &[&CONTENT_TYPE, &CONTENT_LENGTH];
 
         // store original headers, to assert that they are repeated on the 304 response.
         let original_headers: HashMap<HeaderName, HeaderValue> = uncached_response

crates/lib/docs_rs_storage/src/archive_index.rs

@@ -1100,7 +1100,7 @@ mod tests {
                     date_updated: Utc::now(),
                     etag: None,
                     compression: None,
-                    content_length: content.len(),
+                    content_length: Some(content.len()),
                     content: Box::new(Cursor::new(content)),
                 })
             })
@@ -1155,7 +1155,7 @@ mod tests {
                     date_updated: Utc::now(),
                     etag: None,
                     compression: None,
-                    content_length: content.len(),
+                    content_length: Some(content.len()),
                     content: Box::new(Cursor::new(content)),
                 })
             })

crates/lib/docs_rs_storage/src/backends/s3.rs

@@ -461,10 +461,7 @@ impl StorageBackendMethods for S3Backend {
                 .unwrap_or(mime::APPLICATION_OCTET_STREAM),
             date_updated,
             etag,
-            content_length: res
-                .content_length
-                .and_then(|length| length.try_into().ok())
-                .unwrap_or(0),
+            content_length: res.content_length.and_then(|length| length.try_into().ok()),
             content: Box::new(res.body.into_async_read()),
             compression,
         })

crates/lib/docs_rs_storage/src/blob.rs

@@ -88,7 +88,12 @@ pub struct StreamingBlob {
     pub date_updated: DateTime<Utc>,
     pub etag: Option<ETag>,
     pub compression: Option<CompressionAlgorithm>,
-    pub content_length: usize,
+    /// initially the original content length on S3,
+    /// what we return to the user might be bigger because of
+    /// streaming decompression.
+    /// Will be emptied when we decompress the blob because then
+    /// re return more data than we have on S3.
+    pub content_length: Option<usize>,
     pub content: Box<dyn AsyncBufRead + Unpin + Send>,
 }
 
@@ -100,6 +105,7 @@ impl fmt::Debug for StreamingBlob {
             .field("date_updated", &self.date_updated)
             .field("etag", &self.etag)
             .field("compression", &self.compression)
+            .field("content_length", &self.content_length)
             .finish()
     }
 }
@@ -133,14 +139,15 @@ impl StreamingBlob {
         );
 
         self.compression = None;
+        self.content_length = None;
         // not touching the etag, it should represent the original content
         Ok(self)
     }
 
     /// consume the inner stream and materialize the full blob into memory.
     pub async fn materialize(mut self, max_size: usize) -> Result<Blob> {
         let mut content = SizedBuffer::new(max_size);
-        content.reserve(self.content_length);
+        content.reserve(self.content_length.unwrap_or(16 * 1024));
 
         io::copy(&mut self.content, &mut content).await?;
 
@@ -163,7 +170,7 @@ impl From<Blob> for StreamingBlob {
             date_updated: value.date_updated,
             etag: value.etag,
             compression: value.compression,
-            content_length: value.content.len(),
+            content_length: Some(value.content.len()),
             content: Box::new(Cursor::new(value.content)),
         }
     }
@@ -188,7 +195,7 @@ mod test {
             date_updated: Utc::now(),
             compression: alg,
             etag: Some(compute_etag(&content)),
-            content_length: content.len(),
+            content_length: Some(content.len()),
             content: Box::new(Cursor::new(content)),
         }
     }
@@ -200,14 +207,17 @@ mod test {
         // without decompression
         {
             let stream = streaming_blob(CONTENT, None);
+            assert_eq!(stream.content_length, Some(CONTENT.len()));
             let blob = stream.materialize(usize::MAX).await?;
             assert_eq!(blob.content, CONTENT);
             assert!(blob.compression.is_none());
         }
 
         // with decompression, does nothing
         {
-            let stream = streaming_blob(CONTENT, None);
+            let stream = streaming_blob(CONTENT, None).decompress().await?;
+            // no compression, content length stays valid
+            assert_eq!(stream.content_length, Some(CONTENT.len()));
             let blob = stream.decompress().await?.materialize(usize::MAX).await?;
             assert_eq!(blob.content, CONTENT);
             assert!(blob.compression.is_none());
@@ -225,6 +235,7 @@ mod test {
         // Doesn't fail because we don't call `.decompress`
         {
             let stream = streaming_blob(NOT_ZSTD, Some(alg));
+            assert_eq!(stream.content_length, Some(NOT_ZSTD.len()));
             let blob = stream.materialize(usize::MAX).await?;
             assert_eq!(blob.content, NOT_ZSTD);
             assert_eq!(blob.compression, Some(alg));
@@ -267,6 +278,7 @@ mod test {
         // without decompression
         {
             let stream = streaming_blob(compressed_content.clone(), Some(alg));
+            assert_eq!(stream.content_length, Some(compressed_content.len()));
             let blob = stream.materialize(usize::MAX).await?;
             assert_eq!(blob.content, compressed_content);
             assert_eq!(blob.content.last_chunk::<3>().unwrap(), &ZSTD_EOF_BYTES);
@@ -275,11 +287,12 @@ mod test {
 
         // with decompression
         {
-            let blob = streaming_blob(compressed_content.clone(), Some(alg))
+            let stream = streaming_blob(compressed_content.clone(), Some(alg))
                 .decompress()
-                .await?
-                .materialize(usize::MAX)
                 .await?;
+            // content length becomes unknown with decompression
+            assert!(stream.content_length.is_none());
+            let blob = stream.materialize(usize::MAX).await?;
             assert_eq!(blob.content, CONTENT);
             assert!(blob.compression.is_none());
         }

crates/lib/docs_rs_storage/src/compression.rs

@@ -219,19 +219,20 @@ mod tests {
         }
 
         // try decompress via storage API
-        let blob = StreamingBlob {
+        let stream = StreamingBlob {
             path: "some_path.db".into(),
             mime: mime::APPLICATION_OCTET_STREAM,
             date_updated: Utc::now(),
             etag: None,
             compression: Some(alg),
-            content_length: compressed_index_content.len(),
+            content_length: Some(compressed_index_content.len()),
             content: Box::new(io::Cursor::new(compressed_index_content)),
         }
         .decompress()
-        .await?
-        .materialize(usize::MAX)
         .await?;
+        assert!(stream.content_length.is_none());
+
+        let blob = stream.materialize(usize::MAX).await?;
 
         assert_eq!(blob.compression, None);
         assert_eq!(blob.content, CONTENT);