From 9b0154064fb35f4129ae67257c4a630b6fcfb2be Mon Sep 17 00:00:00 2001
From: Jeff Baskin <jeff.l.baskin@gmail.com>
Date: Wed, 15 Feb 2023 07:10:28 -0500
Subject: [PATCH] Handled unicode error.

---
 src/morethantext/cache.rs | 114 +++++++++++++++++++++++++++++++++++++-
 src/morethantext/store.rs |  19 ++++---
 2 files changed, 124 insertions(+), 9 deletions(-)

diff --git a/src/morethantext/cache.rs b/src/morethantext/cache.rs
index 42c11af..77b13a0 100644
--- a/src/morethantext/cache.rs
+++ b/src/morethantext/cache.rs
@@ -1,8 +1,10 @@
-use super::{DBError, SessionData, Store};
+use super::{DBError, FileData, SessionData, Store};
 use async_std::{fs::write, path::Path};
 use rand::{distributions::Alphanumeric, thread_rng, Rng};
 use std::{
     cell::Cell,
+    slice,
+    str,
     time::{Duration, Instant},
 };
 
@@ -40,6 +42,48 @@ impl SessionData for DataType {
     }
 }
 
+impl FileData<Self> for DataType {
+    fn to_bytes(&self) -> Vec<u8> {
+        let mut output = Vec::new();
+        match self {
+            DataType::DBMap(store) => {
+                output.append(&mut "DBMap".as_bytes().to_vec());
+                output.push(0);
+                output.append(&mut store.to_bytes());
+            },
+        }
+        output
+    }
+
+    fn from_bytes(data: &mut slice::Iter<u8>) -> Result<Self, DBError> {
+        let mut header: Vec<u8> = Vec::new();
+        loop {
+            let letter = match data.next() {
+                Some(a) => a.clone(),
+                None => 0,
+            };
+            if letter == 0 {
+                break;
+            } else {
+                header.push(letter);
+            }
+        }
+        let header = match str::from_utf8(&header) {
+            Ok(item) => item,
+            Err(_) => return Err(DBError::new("file corruption")),
+        };
+        match header {
+            "DBMap" => {
+                match Store::from_bytes(data) {
+                    Ok(store) => Ok(DataType::DBMap(store)),
+                    Err(err) => Err(err),
+                }
+            },
+            _ => Err(DBError::new("file corruption")),
+        }
+    }
+}
+
 struct Entry {
     data: DataType,
     filename: String,
@@ -108,7 +152,73 @@ mod datatype_sesssion {
 mod datatype_file {
     use super::*;
 
-    // Test file data traits here.
+    #[test]
+    fn new_store_bytes() {
+        let dbs = DataType::new("store").unwrap();
+        let mut expected = "DBMap".as_bytes().to_vec();
+        expected.push(0);
+        assert_eq!(dbs.to_bytes(), expected);
+    }
+
+    #[test]
+    fn store_bytes_with_info() {
+        let name = "title";
+        let id = "king";
+        let mut store = Store::new();
+        let mut dt_store = DataType::new("store").unwrap();
+        let mut expected = dt_store.to_bytes();
+        store.add("database", name, id);
+        expected.append(&mut store.to_bytes());
+        dt_store.add("database", name, id);
+        assert_eq!(dt_store.to_bytes(), expected);
+    }
+
+    #[test]
+    fn read_empty_store() {
+        let dt_store = DataType::new("store").unwrap();
+        let data = dt_store.to_bytes();
+        let mut feed = data.iter();
+        let output = DataType::from_bytes(&mut feed).unwrap();
+        assert_eq!(dt_store.list(["database"].to_vec()).unwrap(), output.list(["database"].to_vec()).unwrap());
+    }
+
+    #[test]
+    fn read_store_info() {
+        let mut dt_store = DataType::new("store").unwrap();
+        dt_store.add("database", "raven", "beastboy").unwrap();
+        let data = dt_store.to_bytes();
+        let mut feed = data.iter();
+        let output = DataType::from_bytes(&mut feed).unwrap();
+        assert_eq!(dt_store.list(["database"].to_vec()).unwrap(), output.list(["database"].to_vec()).unwrap());
+    }
+
+    #[test]
+    fn read_bad_header() -> Result<(), DBError> {
+        let data = "sdghsdl".as_bytes().to_vec();
+        let mut feed = data.iter();
+        match DataType::from_bytes(&mut feed) {
+            Ok(_) => Err(DBError::new("should have raised an error")),
+            Err(err) => {
+                assert_eq!(err.to_string(), "file corruption");
+                Ok(())
+            },
+        }
+    }
+
+    #[test]
+    fn read_bad_store() -> Result<(), DBError> {
+        let mut data = "DBMap".as_bytes().to_vec();
+        data.push(0);
+        data.append(&mut "sdfgs".as_bytes().to_vec());
+        let mut feed = data.iter();
+        match DataType::from_bytes(&mut feed) {
+            Ok(_) => Err(DBError::new("should have raised an error")),
+            Err(err) => {
+                assert_eq!(err.to_string(), "file corruption");
+                Ok(())
+            },
+        }
+    }
 }
 
 #[cfg(test)]
diff --git a/src/morethantext/store.rs b/src/morethantext/store.rs
index b233718..166c51d 100644
--- a/src/morethantext/store.rs
+++ b/src/morethantext/store.rs
@@ -40,26 +40,31 @@ impl FileData<Self> for Store {
         let mut id: Vec<u8> = Vec::new();
         let mut get_id = false;
         let mut letter: u8;
+        let err_msg = "file corruption";
         loop {
             match data.next() {
                 Some(a) => letter = a.clone(),
                 None => {
                     if !name.is_empty() {
-                        return Err(DBError::new("file corruption"));
+                        return Err(DBError::new(err_msg));
                     }
                     break;
                 }
             }
             if letter == 0 {
                 if get_id {
-                    match output.add(
-                        "database",
-                        str::from_utf8(&name).unwrap(),
-                        str::from_utf8(&id).unwrap(),
-                    ) {
+                    let name_holder = match str::from_utf8(&name) {
+                        Ok(item) => item,
+                        Err(_) => return Err(DBError::new(err_msg)),
+                    };
+                    let id_holder = match str::from_utf8(&id) {
+                        Ok(item) => item,
+                        Err(_) => return Err(DBError::new(err_msg)),
+                    };
+                    match output.add("database", name_holder, id_holder) {
                         Ok(_) => (),
                         Err(err) => {
-                            let mut error = DBError::new("file corruption");
+                            let mut error = DBError::new(err_msg);
                             error.add_source(err);
                             return Err(error);
                         }