Glue for range request ABI

Author: cceckman-at-fastly

lib/src/cache.rs

@@ -263,11 +263,6 @@ impl From<Arc<CacheData>> for Found {
 }
 
 impl Found {
-    /// Access the body of the cached object.
-    pub async fn body(&self) -> Result<Body, crate::Error> {
-        self.get_body().build().await
-    }
-
     fn get_body(&self) -> GetBodyBuilder {
         self.data.as_ref().body()
     }
@@ -616,7 +611,7 @@ mod tests {
 
                 let nonempty = cache.lookup(&key, &HeaderMap::default()).await;
                 let found = nonempty.found().expect("should have found inserted key");
-                let got = found.body().await.unwrap().read_into_vec().await.unwrap();
+                let got = found.get_body().build().await.unwrap().read_into_vec().await.unwrap();
                 assert_eq!(got, value);
             });
         }

lib/src/component/cache.rs

@@ -123,7 +123,7 @@ fn load_lookup_options(
     session: &Session,
     options_mask: api::LookupOptionsMask,
     options: api::LookupOptions,
-) -> Result<HeaderMap, Error> {
+) -> Result<(HeaderMap, bool), Error> {
     let headers = if options_mask.contains(api::LookupOptionsMask::REQUEST_HEADERS) {
         let handle = options.request_headers;
         let parts = session.request_parts(handle.into())?;
@@ -133,10 +133,14 @@ fn load_lookup_options(
     };
     let options_mask = options_mask & !api::LookupOptionsMask::REQUEST_HEADERS;
 
+    let always_use_requested_range =
+        options_mask.contains(api::LookupOptionsMask::ALWAYS_USE_REQUESTED_RANGE);
+    let options_mask = options_mask & !api::LookupOptionsMask::ALWAYS_USE_REQUESTED_RANGE;
+
     if options_mask != api::LookupOptionsMask::empty() {
         return Err(Error::NotAvailable("unknown cache lookup option"));
     }
-    Ok(headers)
+    Ok((headers, always_use_requested_range))
 }
 
 #[async_trait::async_trait]
@@ -147,14 +151,18 @@ impl api::Host for ComponentCtx {
         options_mask: api::LookupOptionsMask,
         options: api::LookupOptions,
     ) -> Result<api::Handle, types::Error> {
-        let headers = load_lookup_options(&self.session, options_mask, options)?;
+        let (headers, always_use_requested_range) =
+            load_lookup_options(&self.session, options_mask, options)?;
 
         let key: CacheKey = get_key(key)?;
         let cache = Arc::clone(self.session.cache());
 
-        let task = PeekableTask::spawn(Box::pin(
-            async move { Ok(cache.lookup(&key, &headers).await) },
-        ))
+        let task = PeekableTask::spawn(Box::pin(async move {
+            Ok(cache
+                .lookup(&key, &headers)
+                .await
+                .with_always_use_requested_range(always_use_requested_range))
+        }))
         .await;
         let task = PendingCacheTask::new(task);
         let handle: CacheHandle = self.session.insert_cache_op(task).into();
@@ -298,33 +306,44 @@ impl api::Host for ComponentCtx {
     async fn get_body(
         &mut self,
         handle: api::Handle,
-        options_mask: api::GetBodyOptionsMask,
-        _options: api::GetBodyOptions,
+        mut options_mask: api::GetBodyOptionsMask,
+        options: api::GetBodyOptions,
     ) -> Result<http_types::BodyHandle, types::Error> {
+        let from = if options_mask.contains(api::GetBodyOptionsMask::FROM) {
+            Some(options.from)
+        } else {
+            None
+        };
+        options_mask &= !api::GetBodyOptionsMask::FROM;
+        let to = if options_mask.contains(api::GetBodyOptionsMask::TO) {
+            Some(options.to)
+        } else {
+            None
+        };
+        options_mask &= !api::GetBodyOptionsMask::TO;
+
         if options_mask != api::GetBodyOptionsMask::empty() {
             return Err(Error::NotAvailable("unknown cache get_body option").into());
         }
 
         // We wind up re-borrowing `found` and `self.session` several times here, to avoid
-        // borrowing the both of them at once. Ultimately it is possible that inserting a body
-        // would change the address of Found, by re-shuffling the AsyncItems table; once again,
-        // borrowck wins the day.
+        // borrowing the both of them at once.
+        // (It possible that inserting a body would change the address of Found, by re-shuffling
+        // the AsyncItems table; we have to live by borrowck's rules.)
         //
         // We have an exclusive borrow &mut self.session for the lifetime of this call,
         // so even though we're re-borrowing/repeating lookups, we know we won't run into TOCTOU.
 
-        let found = self
-            .session
-            .cache_entry(handle.into())
-            .await?
-            .found()
-            .ok_or_else(|| Error::CacheError(cache::Error::Missing))?;
+        let entry = self.session.cache_entry(handle.into()).await?;
 
         // Preemptively (optimistically) start a read. Don't worry, the Drop impl for Body will
         // clean up the copying task.
         // We have to do this to allow `found`'s lifetime to end before self.session.body, which
         // has to re-borrow self.self.session.
-        let body = found.body().await?;
+        let body = entry.body(from, to).await?;
+        let found = entry
+            .found()
+            .ok_or(Error::CacheError(crate::cache::Error::Missing))?;
 
         if let Some(prev_handle) = found.last_body_handle {
             // Check if they're still reading the previous handle.
@@ -364,21 +383,28 @@ impl api::Host for ComponentCtx {
         options_mask: api::LookupOptionsMask,
         options: api::LookupOptions,
     ) -> Result<api::BusyHandle, types::Error> {
-        let headers = load_lookup_options(&self.session, options_mask, options)?;
+        let (headers, always_use_requested_range) =
+            load_lookup_options(&self.session, options_mask, options)?;
 
         let key: CacheKey = get_key(key)?;
         let cache = Arc::clone(self.session.cache());
 
         // Look up once, joining the transaction only if obligated:
-        let e = cache.transaction_lookup(&key, &headers, false).await;
+        let e = cache
+            .transaction_lookup(&key, &headers, false)
+            .await
+            .with_always_use_requested_range(always_use_requested_range);
         let ready = e.found().is_some() || e.go_get().is_some();
         // If we already got _something_, we can provide an already-complete PeekableTask.
         // Otherwise we need to spawn it and let it block in the background.
         let task = if ready {
             PeekableTask::complete(e)
         } else {
             PeekableTask::spawn(Box::pin(async move {
-                Ok(cache.transaction_lookup(&key, &headers, true).await)
+                Ok(cache
+                    .transaction_lookup(&key, &headers, true)
+                    .await
+                    .with_always_use_requested_range(always_use_requested_range))
             }))
             .await
         };

lib/src/wiggle_abi/cache.rs

@@ -139,7 +139,7 @@ fn load_lookup_options(
     memory: &wiggle::GuestMemory<'_>,
     mut options_mask: types::CacheLookupOptionsMask,
     options: wiggle::GuestPtr<types::CacheLookupOptions>,
-) -> Result<HeaderMap, Error> {
+) -> Result<(HeaderMap, bool), Error> {
     let options = memory.read(options)?;
     let headers = if options_mask.contains(types::CacheLookupOptionsMask::REQUEST_HEADERS) {
         let handle = options.request_headers;
@@ -158,11 +158,15 @@ fn load_lookup_options(
         });
     }
 
+    let always_use_requested_range =
+        options_mask.contains(types::CacheLookupOptionsMask::ALWAYS_USE_REQUESTED_RANGE);
+    options_mask &= !types::CacheLookupOptionsMask::ALWAYS_USE_REQUESTED_RANGE;
+
     if !options_mask.is_empty() {
         return Err(Error::NotAvailable("unknown cache lookup option"));
     }
 
-    Ok(headers)
+    Ok((headers, always_use_requested_range))
 }
 
 #[allow(unused_variables)]
@@ -175,13 +179,17 @@ impl FastlyCache for Session {
         options_mask: types::CacheLookupOptionsMask,
         options: wiggle::GuestPtr<types::CacheLookupOptions>,
     ) -> Result<types::CacheHandle, Error> {
-        let headers = load_lookup_options(self, memory, options_mask, options)?;
+        let (headers, always_use_requested_range) =
+            load_lookup_options(self, memory, options_mask, options)?;
         let key = load_cache_key(memory, cache_key)?;
         let cache = Arc::clone(self.cache());
 
-        let task = PeekableTask::spawn(Box::pin(
-            async move { Ok(cache.lookup(&key, &headers).await) },
-        ))
+        let task = PeekableTask::spawn(Box::pin(async move {
+            Ok(cache
+                .lookup(&key, &headers)
+                .await
+                .with_always_use_requested_range(always_use_requested_range))
+        }))
         .await;
         let task = PendingCacheTask::new(task);
         let handle = self.insert_cache_op(task);
@@ -330,20 +338,27 @@ impl FastlyCache for Session {
         options_mask: types::CacheLookupOptionsMask,
         options: wiggle::GuestPtr<types::CacheLookupOptions>,
     ) -> Result<types::CacheBusyHandle, Error> {
-        let headers = load_lookup_options(self, memory, options_mask, options)?;
+        let (headers, always_use_requested_range) =
+            load_lookup_options(self, memory, options_mask, options)?;
         let key = load_cache_key(memory, cache_key)?;
         let cache = Arc::clone(self.cache());
 
         // Look up once, joining the transaction only if obligated:
-        let e = cache.transaction_lookup(&key, &headers, false).await;
+        let e = cache
+            .transaction_lookup(&key, &headers, false)
+            .await
+            .with_always_use_requested_range(always_use_requested_range);
         let ready = e.found().is_some() || e.go_get().is_some();
         // If we already got _something_, we can provide an already-complete PeekableTask.
         // Otherwise we need to spawn it and let it block in the background.
         let task = if ready {
             PeekableTask::complete(e)
         } else {
             PeekableTask::spawn(Box::pin(async move {
-                Ok(cache.transaction_lookup(&key, &headers, true).await)
+                Ok(cache
+                    .transaction_lookup(&key, &headers, true)
+                    .await
+                    .with_always_use_requested_range(always_use_requested_range))
             }))
             .await
         };
@@ -529,31 +544,45 @@ impl FastlyCache for Session {
         &mut self,
         memory: &mut wiggle::GuestMemory<'_>,
         handle: types::CacheHandle,
-        options_mask: types::CacheGetBodyOptionsMask,
+        mut options_mask: types::CacheGetBodyOptionsMask,
         options: &types::CacheGetBodyOptions,
     ) -> Result<types::BodyHandle, Error> {
+        let from = if options_mask.contains(types::CacheGetBodyOptionsMask::FROM) {
+            Some(options.from)
+        } else {
+            None
+        };
+        options_mask &= !types::CacheGetBodyOptionsMask::FROM;
+        let to = if options_mask.contains(types::CacheGetBodyOptionsMask::TO) {
+            Some(options.to)
+        } else {
+            None
+        };
+        options_mask &= !types::CacheGetBodyOptionsMask::TO;
+
         if !options_mask.is_empty() {
             return Err(Error::NotAvailable("unknown cache get_body option").into());
         }
 
         // We wind up re-borrowing `found` and `self.session` several times here, to avoid
-        // borrowing the both of them at once. Ultimately it is possible that inserting a body
-        // would change the address of Found, by re-shuffling the AsyncItems table; once again,
-        // borrowck wins the day.
+        // borrowing the both of them at once.
+        // (It possible that inserting a body would change the address of Found, by re-shuffling
+        // the AsyncItems table; we have to live by borrowck's rules.)
         //
         // We have an exclusive borrow &mut self.session for the lifetime of this call,
         // so even though we're re-borrowing/repeating lookups, we know we won't run into TOCTOU.
 
-        let found = self
-            .cache_entry(handle.into())
-            .await?
-            .found()
-            .ok_or(Error::CacheError(crate::cache::Error::Missing))?;
+        let entry = self.cache_entry(handle.into()).await?;
+
         // Preemptively (optimistically) start a read. Don't worry, the Drop impl for Body will
         // clean up the copying task.
         // We have to do this to allow `found`'s lifetime to end before self.session.body, which
         // has to re-borrow self.self.session.
-        let body = found.body().await?;
+        let body = entry.body(from, to).await?;
+
+        let found = entry
+            .found()
+            .ok_or(Error::CacheError(crate::cache::Error::Missing))?;
 
         if let Some(prev_handle) = found.last_body_handle {
             // Check if they're still reading the previous handle.