From ae6cfd3baf659600cecf746a57fb89a564279ca3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20=C5=A0pa=C4=8Dek?= <jan.spacek@keyrock.eu>
Date: Thu, 18 Jan 2024 10:15:50 +0100
Subject: [PATCH] Improve optional support for borsh serialization

- Do not allocate when deserializing ArrayString
- Serialize length as u32, not as u64, to be consistent with
  serialization of [T] and str
- Add tests
---
 .github/workflows/ci.yml |  4 +--
 src/array_string.rs      | 21 ++++++++----
 src/arrayvec.rs          | 11 ++----
 tests/borsh.rs           | 73 ++++++++++++++++++++++++++++++++++++++++
 4 files changed, 93 insertions(+), 16 deletions(-)
 create mode 100644 tests/borsh.rs

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 56fdb0f..6f8c77b 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -18,7 +18,7 @@ jobs:
       matrix:
         include:
           - rust: 1.51.0 # MSRV
-            features: serde
+            features: serde, borsh
             experimental: false
           - rust: stable
             features:
@@ -28,7 +28,7 @@ jobs:
             features: serde
             experimental: false
           - rust: nightly
-            features: serde, zeroize
+            features: serde, borsh, zeroize
             experimental: false
 
     steps:
diff --git a/src/array_string.rs b/src/array_string.rs
index 7145427..1144360 100644
--- a/src/array_string.rs
+++ b/src/array_string.rs
@@ -640,13 +640,22 @@ impl<const CAP: usize> borsh::BorshSerialize for ArrayString<CAP> {
 /// Requires crate feature `"borsh"`
 impl<const CAP: usize> borsh::BorshDeserialize for ArrayString<CAP> {
     fn deserialize_reader<R: borsh::io::Read>(reader: &mut R) -> borsh::io::Result<Self> {
-        let s = <String as borsh::BorshDeserialize>::deserialize_reader(reader)?;
-        ArrayString::from(&s).map_err(|_| {
-            borsh::io::Error::new(
+        let len = <u32 as borsh::BorshDeserialize>::deserialize_reader(reader)? as usize;
+        if len > CAP {
+            return Err(borsh::io::Error::new(
                 borsh::io::ErrorKind::InvalidData,
-                format!("expected a string no more than {} bytes long", CAP),
-            )
-        })
+                format!("Expected a string no more than {} bytes long", CAP),
+            ))
+        }
+
+        let mut buf = [0u8; CAP];
+        let buf = &mut buf[..len];
+        reader.read_exact(buf)?;
+
+        let s = str::from_utf8(&buf).map_err(|err| {
+            borsh::io::Error::new(borsh::io::ErrorKind::InvalidData, err.to_string())
+        })?;
+        Ok(Self::from(s).unwrap())
     }
 }
 
diff --git a/src/arrayvec.rs b/src/arrayvec.rs
index 9734dd2..893028b 100644
--- a/src/arrayvec.rs
+++ b/src/arrayvec.rs
@@ -1309,12 +1309,7 @@ where
     T: borsh::BorshSerialize,
 {
     fn serialize<W: borsh::io::Write>(&self, writer: &mut W) -> borsh::io::Result<()> {
-        let vs = self.as_slice();
-        <usize as borsh::BorshSerialize>::serialize(&vs.len(), writer)?;
-        for elem in vs {
-            <T as borsh::BorshSerialize>::serialize(elem, writer)?;
-        }
-        Ok(())
+        <[T] as borsh::BorshSerialize>::serialize(self.as_slice(), writer)
     }
 }
 
@@ -1326,13 +1321,13 @@ where
 {
     fn deserialize_reader<R: borsh::io::Read>(reader: &mut R) -> borsh::io::Result<Self> {
         let mut values = Self::new();
-        let len = <usize as borsh::BorshDeserialize>::deserialize_reader(reader)?;
+        let len = <u32 as borsh::BorshDeserialize>::deserialize_reader(reader)?;
         for _ in 0..len {
             let elem = <T as borsh::BorshDeserialize>::deserialize_reader(reader)?;
             if let Err(_) = values.try_push(elem) {
                 return Err(borsh::io::Error::new(
                     borsh::io::ErrorKind::InvalidData,
-                    format!("expected an array with no more than {} items", CAP),
+                    format!("Expected an array with no more than {} items", CAP),
                 ));
             }
         }
diff --git a/tests/borsh.rs b/tests/borsh.rs
new file mode 100644
index 0000000..f05abcf
--- /dev/null
+++ b/tests/borsh.rs
@@ -0,0 +1,73 @@
+#![cfg(feature = "borsh")]
+use std::fmt;
+extern crate arrayvec;
+extern crate borsh;
+
+fn assert_ser<T: borsh::BorshSerialize>(v: &T, expected_bytes: &[u8]) {
+    let mut actual_bytes = Vec::new();
+    v.serialize(&mut actual_bytes).unwrap();
+    assert_eq!(actual_bytes, expected_bytes);
+}
+
+fn assert_roundtrip<T: borsh::BorshSerialize + borsh::BorshDeserialize + PartialEq + fmt::Debug>(v: &T) {
+    let mut bytes = Vec::new();
+    v.serialize(&mut bytes).unwrap();
+    let v_de = T::try_from_slice(&bytes).unwrap();
+    assert_eq!(*v, v_de);
+}
+
+mod array_vec {
+    use arrayvec::ArrayVec;
+    use super::{assert_ser, assert_roundtrip};
+
+    #[test]
+    fn test_empty() {
+        let vec = ArrayVec::<u32, 0>::new();
+        assert_ser(&vec, b"\0\0\0\0");
+        assert_roundtrip(&vec);
+    }
+
+    #[test]
+    fn test_full() {
+        let mut vec = ArrayVec::<u32, 3>::new();
+        vec.push(0xdeadbeef);
+        vec.push(0x123);
+        vec.push(0x456);
+        assert_ser(&vec, b"\x03\0\0\0\xef\xbe\xad\xde\x23\x01\0\0\x56\x04\0\0");
+        assert_roundtrip(&vec);
+    }
+
+    #[test]
+    fn test_with_free_capacity() {
+        let mut vec = ArrayVec::<u32, 3>::new();
+        vec.push(0xdeadbeef);
+        assert_ser(&vec, b"\x01\0\0\0\xef\xbe\xad\xde");
+        assert_roundtrip(&vec);
+    }
+}
+
+mod array_string {
+    use arrayvec::ArrayString;
+    use super::{assert_ser, assert_roundtrip};
+
+    #[test]
+    fn test_empty() {
+        let string = ArrayString::<0>::new();
+        assert_ser(&string, b"\0\0\0\0");
+        assert_roundtrip(&string);
+    }
+
+    #[test]
+    fn test_full() {
+        let string = ArrayString::from_byte_string(b"hello world").unwrap();
+        assert_ser(&string, b"\x0b\0\0\0hello world");
+        assert_roundtrip(&string);
+    }
+
+    #[test]
+    fn test_with_free_capacity() {
+        let string = ArrayString::<16>::from("hello world").unwrap();
+        assert_ser(&string, b"\x0b\0\0\0hello world");
+        assert_roundtrip(&string);
+    }
+}