From 9d57775b62386552215154f6c484a2fd31f33139 Mon Sep 17 00:00:00 2001
From: Miraculous Owonubi <omiraculous@gmail.com>
Date: Wed, 31 May 2023 16:03:39 +0100
Subject: [PATCH] chore: fix panic on arbitrary `AccountId` generation (#9124)

Fixes unintended panic in the `Arbitrary` implementation of `AccountId`.

As pointed out in https://github.com/near/nearcore/pull/9107#discussion_r1207737121, certain `Unstructured` data sources like `a|bcd` or `a_-b` causes a panic with the current implementation.

This patch reduces the consumed input to the limit of validity adherence. Failing otherwise. But never panics.

Worth noting this implementation does not provide a fallback for input sources shorter than `AccountId::MIN_LEN` or longer than `AccountId::MAX_LEN`.
---
 core/account-id/src/lib.rs | 41 ++++++++++++++++++++++++++++----------
 1 file changed, 30 insertions(+), 11 deletions(-)

diff --git a/core/account-id/src/lib.rs b/core/account-id/src/lib.rs
index 0776261231b..196232f03b8 100644
--- a/core/account-id/src/lib.rs
+++ b/core/account-id/src/lib.rs
@@ -377,18 +377,15 @@ impl<'a> arbitrary::Arbitrary<'a> for AccountId {
     }
 
     fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
-        let s = u.arbitrary::<&str>()?;
-        match s.parse::<AccountId>() {
-            Ok(account_id) => Ok(account_id),
-            Err(e) => {
-                if let Some((valid_up_to, _)) = e.char {
-                    let valid = &s[..valid_up_to];
-                    debug_assert!(AccountId::validate(valid).is_ok());
-                    #[allow(deprecated)]
-                    Ok(AccountId::new_unvalidated(valid.to_string()))
-                } else {
-                    Err(arbitrary::Error::IncorrectFormat)
+        let mut s = u.arbitrary::<&str>()?;
+        loop {
+            match s.parse::<AccountId>() {
+                Ok(account_id) => break Ok(account_id),
+                Err(ParseAccountError { char: Some((idx, _)), .. }) => {
+                    s = &s[..idx];
+                    continue;
                 }
+                _ => break Err(arbitrary::Error::IncorrectFormat),
             }
         }
     }
@@ -722,4 +719,26 @@ mod tests {
             );
         }
     }
+
+    #[test]
+    #[cfg(feature = "arbitrary")]
+    fn test_arbitrary() {
+        let corpus = [
+            ("a|bcd", None),
+            ("ab|cde", Some("ab")),
+            ("a_-b", None),
+            ("ab_-c", Some("ab")),
+            ("a", None),
+            ("miraclx.near", Some("miraclx.near")),
+            ("01234567890123456789012345678901234567890123456789012345678901234", None),
+        ];
+
+        for (input, expected_output) in corpus {
+            assert!(input.len() <= u8::MAX as usize);
+            let data = [input.as_bytes(), &[input.len() as _]].concat();
+            let mut u = arbitrary::Unstructured::new(&data);
+
+            assert_eq!(u.arbitrary::<AccountId>().as_deref().ok(), expected_output);
+        }
+    }
 }