Fix header_values_set handling in the component implementation

fastly · May 20, 2024 · d7be5fa · d7be5fa
1 parent 7040716
commit d7be5fa
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 16 deletions.
diff --git a/lib/src/component/http_req.rs b/lib/src/component/http_req.rs
@@ -257,15 +257,22 @@ impl http_req::Host for Session {
         let headers = &mut self.request_parts_mut(h.into())?.headers;
 
         let name = HeaderName::from_bytes(name.as_bytes())?;
+        let values = {
+            // split slice along nul bytes
+            let mut iter = values.split(|b| *b == 0);
+            // drop the empty item at the end
+            iter.next_back();
+            iter.map(HeaderValue::from_bytes)
+                .collect::<Result<Vec<HeaderValue>, _>>()?
+        };
 
         // Remove any values if they exist
         if let http::header::Entry::Occupied(e) = headers.entry(&name) {
             e.remove_entry_mult();
         }
 
-        // Add all the new values
-        for value in values.split(|b| *b == 0) {
-            headers.append(&name, HeaderValue::from_bytes(value)?);
+        for value in values {
+            headers.append(&name, value);
         }
 
         Ok(())

diff --git a/lib/src/component/http_resp.rs b/lib/src/component/http_resp.rs
@@ -161,15 +161,22 @@ impl http_resp::Host for Session {
         let headers = &mut self.response_parts_mut(h.into())?.headers;
 
         let name = HeaderName::from_bytes(name.as_bytes())?;
+        let values = {
+            // split slice along nul bytes
+            let mut iter = values.split(|b| *b == 0);
+            // drop the empty item at the end
+            iter.next_back();
+            iter.map(HeaderValue::from_bytes)
+                .collect::<Result<Vec<HeaderValue>, _>>()?
+        };
 
         // Remove any values if they exist
         if let http::header::Entry::Occupied(e) = headers.entry(&name) {
             e.remove_entry_mult();
         }
 
-        // Add all the new values
-        for value in values.split(|b| *b == 0) {
-            headers.append(&name, HeaderValue::from_bytes(value)?);
+        for value in values {
+            headers.append(&name, value);
         }
 
         Ok(())

diff --git a/lib/src/wiggle_abi/headers.rs b/lib/src/wiggle_abi/headers.rs
@@ -118,17 +118,15 @@ impl HttpHeaders for HeaderMap<HeaderValue> {
         }
 
         let name = HeaderName::from_bytes(&name.as_slice()?.ok_or(Error::SharedMemory)?)?;
-        let values = values
-            .as_slice()?
-            .ok_or(Error::SharedMemory)?
+        let values = {
+            let values_bytes = values.as_slice()?.ok_or(Error::SharedMemory)?;
             // split slice along nul bytes
-            .split(|b| *b == 0)
-            // reverse and skip to drop the empty item at the end
-            .rev()
-            .skip(1)
-            .map(HeaderValue::from_bytes)
-            // Collect here in order to return early in error case, before modifying headers state
-            .collect::<Result<Vec<HeaderValue>, _>>()?;
+            let mut iter = values_bytes.split(|b| *b == 0);
+            // drop the empty item at the end
+            iter.next_back();
+            iter.map(HeaderValue::from_bytes)
+                .collect::<Result<Vec<HeaderValue>, _>>()?
+        };
 
         // Remove any values if they exist
         if let http::header::Entry::Occupied(e) = self.entry(&name) {