Correct substitute behavior on reallocation and add tests

This fixes the following issue in replace_impl: the call to
try_reserve passed in a difference of capacities, but try_reserve
expects a difference between the desired capacity and the length.
Because the initial capacity was nonzero but the length was zero, this
caused us to reserve less capacity than we ought to have, leading to an
OOB write.

Fix this by reworking replace_impl to have less unsafe code. Now we zero
initialize the buffer, but we also prefer a stack buffer so we may save
an allocation - probably a wash overall.

Add a test for this case.

Author: ridiculousfish

src/ffi.rs

@@ -90,7 +90,7 @@ pub trait CodeUnitWidth: std::fmt::Debug {
     type pcre2_match_context;
     type pcre2_match_data;
     type pcre2_jit_stack;
-    type PCRE2_CHAR: TryInto<Self::SubjectChar>;
+    type PCRE2_CHAR: Default + Copy + TryInto<Self::SubjectChar>;
     type PCRE2_SPTR;
     type name_table_entry: NameTableEntry;
     type SubjectChar: Copy;

src/regex_impl.rs

@@ -546,40 +546,47 @@ impl<W: CodeUnitWidth> Regex<W> {
             options |= PCRE2_SUBSTITUTE_GLOBAL;
         }
 
-        // TODO: we can use MaybeUninit to avoid allocation
-        let mut capacity = 256;
-        let mut output: Vec<W::PCRE2_CHAR> = Vec::with_capacity(capacity);
-        capacity = output.capacity();
+        // We prefer to allocate on the stack but fall back to the heap.
+        // Note that PCRE2 has the following behavior with PCRE2_SUBSTITUTE_OVERFLOW_LENGTH:
+        //   - We supply the initial output buffer size in `capacity`. This should have sufficient
+        //     capacity for the terminating NUL character.
+        //   - If the capacity is NOT sufficient, PCRE2 returns the new required capacity, also
+        //     including the terminating NUL character.
+        //   - If the capacity IS sufficient, PCRE2 returns the number of characters written, NOT
+        //     including the terminating NUL character.
+        // Example: our initial capacity is 256. If the returned string needs to be of length 512,
+        // then PCRE2 will report NOMEMORY and set capacity to 513. After reallocating we pass in
+        // a capacity of 513; it succeeds and sets capacity to 512, which is the length of the result.
+        let mut stack_storage: [W::PCRE2_CHAR; 256] = [W::PCRE2_CHAR::default(); 256];
+        let mut heap_storage = Vec::new();
+        let mut output = stack_storage.as_mut();
+        let mut capacity = output.len();
 
         let mut rc = unsafe {
             self.code
-                .substitute(subject, replacement, 0, options, &mut output, &mut capacity)
+                .substitute(subject, replacement, 0, options, output, &mut capacity)
         };
 
         if let Err(e) = &rc {
             if e.code() == PCRE2_ERROR_NOMEMORY {
-                if output.try_reserve(capacity - output.capacity()).is_err() {
+                if heap_storage.try_reserve_exact(capacity).is_err() {
                     return Err(rc.unwrap_err());
                 }
-                capacity = output.capacity();
+                heap_storage.resize(capacity, W::PCRE2_CHAR::default());
+                output = &mut heap_storage;
+                capacity = output.len();
                 rc = unsafe {
-                    self.code.substitute(
-                        subject,
-                        replacement,
-                        0,
-                        options,
-                        &mut output,
-                        &mut capacity,
-                    )
+                    self.code
+                        .substitute(subject, replacement, 0, options, output, &mut capacity)
                 };
             }
         }
 
         let s = match rc? {
             0 => Cow::Borrowed(subject),
             _ => {
-                // +1 to account for null terminator
-                unsafe { output.set_len(capacity + 1) };
+                // capacity has been updated with the length of the result (excluding nul terminator).
+                let output = &output[..capacity];
 
                 // All inputs contained valid chars, so we expect all outputs to as well.
                 let to_char = |c: W::PCRE2_CHAR| -> W::SubjectChar {
@@ -588,13 +595,7 @@ impl<W: CodeUnitWidth> Regex<W> {
                 };
 
                 // this is really just a type cast
-                let x: Vec<W::SubjectChar> = output
-                    .into_iter()
-                    .map(to_char)
-                    // we don't want to return the null terminator
-                    .take(capacity)
-                    .collect::<Vec<W::SubjectChar>>();
-
+                let x: Vec<W::SubjectChar> = output.iter().copied().map(to_char).collect();
                 Cow::Owned(x)
             }
         };

src/utf32.rs

@@ -130,6 +130,145 @@ mod tests {
         assert_eq!(replaced, &*b("bc"));
     }
 
+    #[test]
+    fn replace_no_match() {
+        let re = RegexBuilder::new().build(b("d")).unwrap();
+        let s = b("abc");
+        let r = b("");
+        let replaced = re.replace(&s, &r, true).unwrap();
+        assert!(
+            matches!(replaced, Cow::Borrowed(_)),
+            "when there is no match, the original string should be returned"
+        );
+        let replaced = replaced.into_owned();
+        assert_eq!(replaced, &*b("abc"));
+    }
+
+    #[test]
+    fn replace_with_replacement() {
+        let re = RegexBuilder::new().build(b("b")).unwrap();
+        let s = b("abc");
+        let r = b("d");
+        let replaced = re.replace(&s, &r, true).unwrap();
+        assert!(
+            matches!(replaced, Cow::Owned(_)),
+            "a replacement should give a new string"
+        );
+        let replaced = replaced.into_owned();
+        assert_eq!(replaced, &*b("adc"));
+    }
+
+    #[test]
+    fn replace_first_occurrence() {
+        let re = RegexBuilder::new().build(b("a")).unwrap();
+        let s = b("aaa");
+        let r = b("b");
+        let replaced = re.replace(&s, &r, false).unwrap();
+        assert!(
+            matches!(replaced, Cow::Owned(_)),
+            "a replacement should give a new string"
+        );
+        let replaced = replaced.into_owned();
+        assert_eq!(replaced, &*b("baa"));
+    }
+
+    #[test]
+    fn replace_multiple_occurrences() {
+        let re = RegexBuilder::new().build(b("a")).unwrap();
+        let s = b("aaa");
+        let r = b("b");
+        let replaced = re.replace_all(&s, &r, false).unwrap();
+        assert!(
+            matches!(replaced, Cow::Owned(_)),
+            "a replacement should give a new string"
+        );
+        let replaced = replaced.into_owned();
+        assert_eq!(replaced, &*b("bbb"));
+    }
+
+    #[test]
+    fn replace_empty_string() {
+        let re = RegexBuilder::new().build(b("")).unwrap();
+        let s = b("abc");
+        let r = b("d");
+        let replaced = re.replace(&s, &r, true).unwrap();
+        assert!(
+            matches!(replaced, Cow::Owned(_)),
+            "a replacement should give a new string"
+        );
+        let replaced = replaced.into_owned();
+        assert_eq!(replaced, &*b("dabc"));
+    }
+
+    #[test]
+    fn replace_empty_with_empty() {
+        let re = RegexBuilder::new().build(b("")).unwrap();
+        let s = b("");
+        let r = b("");
+        let replaced = re.replace(&s, &r, true).unwrap().into_owned();
+        assert_eq!(replaced, &*b(""));
+    }
+
+    #[test]
+    fn replace_long_string() {
+        let long_string = vec!['a'; 1024]; // Create a 1MB string filled with 'a'
+        let re = RegexBuilder::new().build(b("a")).unwrap();
+        let r = b("b");
+        let replaced = re.replace(&long_string, &r, false).unwrap();
+        assert!(
+            matches!(replaced, Cow::Owned(_)),
+            "a replacement should give a new string"
+        );
+        let replaced = replaced.into_owned();
+        let mut expected = long_string.clone();
+        expected[0] = 'b';
+        assert_eq!(replaced, expected);
+    }
+
+    #[test]
+    fn replace_long_string_all() {
+        let long_string = vec!['a'; 1024];
+        let re = RegexBuilder::new().build(b("a")).unwrap();
+        let r = b("b");
+        let replaced = re.replace_all(&long_string, &r, false).unwrap();
+        assert!(
+            matches!(replaced, Cow::Owned(_)),
+            "a replacement should give a new string"
+        );
+        let replaced = replaced.into_owned();
+        let all_b = vec!['b'; 1024];
+        assert_eq!(replaced, all_b);
+    }
+
+    #[test]
+    fn replace_long_string_all_elongating() {
+        let long_string = vec!['a'; 1024];
+        let re = RegexBuilder::new().build(b("a")).unwrap();
+        let r = b("bx");
+        let replaced = re.replace_all(&long_string, &r, false).unwrap();
+        assert!(
+            matches!(replaced, Cow::Owned(_)),
+            "a replacement should give a new string"
+        );
+        let replaced = replaced.into_owned();
+        let mut all_bx = Vec::new();
+        for _ in long_string {
+            all_bx.push('b');
+            all_bx.push('x');
+        }
+        assert_eq!(replaced, all_bx);
+    }
+
+    #[test]
+    fn replace_long_string_all_disappearing() {
+        let long_string = vec!['a'; 1024];
+        let re = RegexBuilder::new().build(b("a")).unwrap();
+        let r = b("");
+        let replaced = re.replace_all(&long_string, &r, false).unwrap();
+        let replaced = replaced.into_owned();
+        assert_eq!(replaced, &[]);
+    }
+
     #[test]
     fn ucp() {
         let re = RegexBuilder::new().ucp(false).build(b(r"\w")).unwrap();