diff --git a/src/db.rs b/src/db.rs index b6b67af..e68f2d3 100644 --- a/src/db.rs +++ b/src/db.rs @@ -82,9 +82,12 @@ impl CommentId { } pub fn from_base64(s: &str) -> Result { - // TODO prevent panic when s is too long let mut buf = [0; 16]; - base64::decode_config_slice(s, base64::URL_SAFE_NO_PAD, &mut buf).map(|_| Self(buf)) + std::panic::catch_unwind(move || { + base64::decode_config_slice(s, base64::URL_SAFE_NO_PAD, &mut buf) + }) + .map_err(|_| base64::DecodeError::InvalidLength)? + .map(|_| Self(buf)) } } @@ -96,6 +99,8 @@ impl AsRef<[u8]> for CommentId { #[cfg(test)] mod test { + use super::*; + #[test] fn test_typed_sled() { let db = sled::Config::new().temporary(true).open().unwrap(); @@ -106,4 +111,9 @@ mod test { //let mut iter = tree.iter(); assert_eq!(iter.next(), Some(Ok(((123, 456), ())))); } + + #[test] + fn test_from_base64_dont_panic() { + assert_eq!(CommentId::from_base64("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"), Err(base64::DecodeError::InvalidLength)); + } }