From b2a6ba6ef6313baa0df137194235d02db6507193 Mon Sep 17 00:00:00 2001
From: Thomas <denhollander.thomas@gmail.com>
Date: Thu, 6 Feb 2020 15:26:59 +0100
Subject: [PATCH] api: add Arbitrary impl for CString

We add a little sophistication here for whether the CString is
completely valid UTF-8 or whether it's just an arbitrary mix of bytes.
(Excluding NUL of course.)

Fixes #165, Closes #257
---
 src/arbitrary.rs | 50 +++++++++++++++++++++++++++++++++++++++++-------
 src/tests.rs     |  5 +++++
 2 files changed, 48 insertions(+), 7 deletions(-)

diff --git a/src/arbitrary.rs b/src/arbitrary.rs
index a1db83a..9333f26 100644
--- a/src/arbitrary.rs
+++ b/src/arbitrary.rs
@@ -3,7 +3,7 @@ use std::collections::{
     BTreeMap, BTreeSet, BinaryHeap, HashMap, HashSet, LinkedList, VecDeque,
 };
 use std::env;
-use std::ffi::OsString;
+use std::ffi::{CString, OsString};
 use std::hash::{BuildHasher, Hash};
 use std::iter::{empty, once};
 use std::net::{
@@ -277,7 +277,7 @@ impl<A: Arbitrary> Arbitrary for Vec<A> {
             let s = g.size();
             g.gen_range(0, s)
         };
-        (0..size).map(|_| Arbitrary::arbitrary(g)).collect()
+        (0..size).map(|_| A::arbitrary(g)).collect()
     }
 
     fn shrink(&self) -> Box<dyn Iterator<Item = Vec<A>>> {
@@ -598,11 +598,7 @@ impl Arbitrary for String {
             let s = g.size();
             g.gen_range(0, s)
         };
-        let mut s = String::with_capacity(size);
-        for _ in 0..size {
-            s.push(char::arbitrary(g));
-        }
-        s
+        (0..size).map(|_| char::arbitrary(g)).collect()
     }
 
     fn shrink(&self) -> Box<dyn Iterator<Item = String>> {
@@ -612,6 +608,46 @@ impl Arbitrary for String {
     }
 }
 
+impl Arbitrary for CString {
+    fn arbitrary<G: Gen>(g: &mut G) -> Self {
+        let size = {
+            let s = g.size();
+            g.gen_range(0, s)
+        };
+        // Use either random bytes or random UTF-8 encoded codepoints.
+        let utf8: bool = g.gen();
+        if utf8 {
+            CString::new(
+                (0..)
+                    .map(|_| char::arbitrary(g))
+                    .filter(|&c| c != '\0')
+                    .take(size)
+                    .collect::<String>(),
+            )
+        } else {
+            CString::new(
+                (0..)
+                    .map(|_| u8::arbitrary(g))
+                    .filter(|&c| c != b'\0')
+                    .take(size)
+                    .collect::<Vec<u8>>(),
+            )
+        }
+        .expect("null characters should have been filtered out")
+    }
+
+    fn shrink(&self) -> Box<dyn Iterator<Item = CString>> {
+        // Use the implementation for a vec here, but make sure null characters
+        // are filtered out.
+        Box::new(VecShrinker::new(self.as_bytes().to_vec()).map(|bytes| {
+            CString::new(
+                bytes.into_iter().filter(|&c| c != 0).collect::<Vec<u8>>(),
+            )
+            .expect("null characters should have been filtered out")
+        }))
+    }
+}
+
 impl Arbitrary for char {
     fn arbitrary<G: Gen>(g: &mut G) -> char {
         let mode = g.gen_range(0, 100);
diff --git a/src/tests.rs b/src/tests.rs
index 7367432..5a02a3f 100644
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -1,6 +1,7 @@
 use std::cmp::Ord;
 use std::collections::hash_map::DefaultHasher;
 use std::collections::{HashMap, HashSet};
+use std::ffi::CString;
 use std::hash::BuildHasherDefault;
 use std::path::PathBuf;
 
@@ -286,4 +287,8 @@ quickcheck! {
     ) -> bool {
         true
     }
+
+    fn cstring(_p: CString) -> bool {
+        true
+    }
 }