From 4df813a1158b8270fc940de95504565f49733d85 Mon Sep 17 00:00:00 2001
From: Tomer Filiba <tomer@sweet.security>
Date: Sat, 10 Aug 2024 16:16:01 +0300
Subject: [PATCH] Crasher: test collection inserts

---
 vicky-crasher/src/main.rs | 109 +++++++++++++++++++++++++++++---------
 1 file changed, 84 insertions(+), 25 deletions(-)

diff --git a/vicky-crasher/src/main.rs b/vicky-crasher/src/main.rs
index 0f5afec..461ed1e 100644
--- a/vicky-crasher/src/main.rs
+++ b/vicky-crasher/src/main.rs
@@ -35,7 +35,7 @@ fn child_inserts() -> Result<()> {
 }
 
 fn child_removals() -> Result<()> {
-    // our job is to create 1M entries while being killed by our evil parent
+    // our job is to remove 1M entries while being killed by our evil parent
 
     let store = VickyStore::open("dbdir", Config::default())?;
     let lowest_bytes = store.get("lowest")?.unwrap_or(vec![0, 0, 0, 0]);
@@ -62,13 +62,43 @@ fn child_removals() -> Result<()> {
     Ok(())
 }
 
+fn child_collection_inserts() -> Result<()> {
+    // our job is to insert 1M entries to a collection while being killed by our evil parent
+
+    let store = VickyStore::open("dbdir", Config::default())?;
+
+    let highest_bytes = store.get("coll_highest")?.unwrap_or(vec![0, 0, 0, 0]);
+    let highest = u32::from_le_bytes([
+        highest_bytes[0],
+        highest_bytes[1],
+        highest_bytes[2],
+        highest_bytes[3],
+    ]);
+
+    if highest == TARGET - 1 {
+        println!("child finished (already at {highest})");
+        return Ok(());
+    }
+
+    println!("child starting at {highest}");
+
+    for i in highest..TARGET {
+        store.set_in_collection("xxx", &i.to_le_bytes(), "yyy")?;
+        store.set("coll_highest", &i.to_le_bytes())?;
+    }
+    println!("child finished");
+
+    Ok(())
+}
+
 fn parent_run(mut child_func: impl FnMut() -> Result<()>, sleep: Range<u64>) -> Result<()> {
     for i in 0.. {
         let pid = unsafe { libc::fork() };
         assert!(pid >= 0);
         if pid == 0 {
             let res = child_func();
-            unsafe { libc::exit(if res.is_err() { 1 } else { 0 }) };
+            res.unwrap();
+            unsafe { libc::exit(0) };
         } else {
             // parent
             std::thread::sleep(Duration::from_millis(
@@ -84,6 +114,10 @@ fn parent_run(mut child_func: impl FnMut() -> Result<()>, sleep: Range<u64>) ->
                 };
             } else {
                 assert!(rc > 0);
+                if !libc::WIFSIGNALED(status) && libc::WEXITSTATUS(status) != 0 {
+                    panic!("child crashed at iteration {i}")
+                }
+
                 println!("child finished in {i} iterations");
                 break;
             }
@@ -97,36 +131,61 @@ fn main() -> Result<()> {
 
     parent_run(child_inserts, 10..300)?;
 
-    println!("Parent starts validating the DB...");
+    {
+        println!("Parent starts validating the DB...");
+
+        let store = VickyStore::open("dbdir", Config::default())?;
+        assert_eq!(
+            store.remove("highest")?,
+            Some((TARGET - 1).to_le_bytes().to_vec())
+        );
+        let mut count = 0;
+        for res in store.iter() {
+            let (k, v) = res?;
+            assert_eq!(v, b"i am a key");
+            let k = u32::from_le_bytes([k[0], k[1], k[2], k[3]]);
+            assert!(k < TARGET);
+            count += 1;
+        }
+        assert_eq!(count, TARGET);
 
-    let store = VickyStore::open("dbdir", Config::default())?;
-    assert_eq!(
-        store.remove("highest")?,
-        Some((TARGET - 1).to_le_bytes().to_vec())
-    );
-    let mut count = 0;
-    for res in store.iter() {
-        let (k, v) = res?;
-        assert_eq!(v, b"i am a key");
-        let k = u32::from_le_bytes([k[0], k[1], k[2], k[3]]);
-        assert!(k < TARGET);
-        count += 1;
+        println!("DB validated successfully");
     }
-    assert_eq!(count, TARGET);
-
-    println!("DB validated successfully");
 
     parent_run(child_removals, 10..30)?;
 
-    println!("Parent starts validating the DB...");
+    {
+        println!("Parent starts validating the DB...");
+
+        let store = VickyStore::open("dbdir", Config::default())?;
+        assert_eq!(
+            store.remove("lowest")?,
+            Some((TARGET - 1).to_le_bytes().to_vec())
+        );
+        assert_eq!(store.iter().count(), 0);
+
+        println!("DB validated successfully");
+    }
+
+    parent_run(child_collection_inserts, 10..30)?;
 
-    assert_eq!(
-        store.remove("lowest")?,
-        Some((TARGET - 1).to_le_bytes().to_vec())
-    );
-    assert_eq!(store.iter().count(), 0);
+    {
+        println!("Parent starts validating the DB...");
 
-    println!("DB validated successfully");
+        let store = VickyStore::open("dbdir", Config::default())?;
+        assert_eq!(
+            store.remove("coll_highest")?,
+            Some((TARGET - 1).to_le_bytes().to_vec())
+        );
+
+        for (i, res) in store.iter_collection("xxx").enumerate() {
+            let (k, v) = res?.unwrap();
+            assert_eq!(k, (i as u32).to_le_bytes());
+            assert_eq!(v, b"yyy");
+        }
+
+        println!("DB validated successfully");
+    }
 
     Ok(())
 }