From 55d1c9532f9a38508aead88ee78909b38880d2a3 Mon Sep 17 00:00:00 2001
From: "Ilya (Marshal)" <ilya@marshal.dev>
Date: Fri, 20 Jun 2025 19:31:58 +0200
Subject: [PATCH] Fix panic on encoding map with integer keys

---
 pytests/test_dag_cbor.py | 16 ++++++++++++++++
 src/lib.rs               | 20 +++++++++++++-------
 2 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/pytests/test_dag_cbor.py b/pytests/test_dag_cbor.py
index f099e3f..b5d1dd1 100644
--- a/pytests/test_dag_cbor.py
+++ b/pytests/test_dag_cbor.py
@@ -135,3 +135,19 @@ def test_recursion_limit_exceed_on_nested_maps() -> None:
         libipld.decode_dag_cbor(dag_cbor)
 
     assert 'in DAG-CBOR decoding' in str(exc_info.value)
+
+
+def test_dab_cbor_decode_map_int_key() -> None:
+    dag_cbor = bytes.fromhex('a10000')
+    with pytest.raises(ValueError) as exc_info:
+        libipld.decode_dag_cbor(dag_cbor)
+
+    assert 'Map keys must be strings' in str(exc_info.value)
+
+
+def test_dab_cbor_encode_map_int_key() -> None:
+    obj = {0: 'value'}
+    with pytest.raises(ValueError) as exc_info:
+        libipld.encode_dag_cbor(obj)
+
+    assert 'Map keys must be strings' in str(exc_info.value)
diff --git a/src/lib.rs b/src/lib.rs
index 290eb2d..483c906 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -51,18 +51,24 @@ fn map_key_cmp(a: &Vec<u8>, b: &Vec<u8>) -> std::cmp::Ordering {
     }
 }
 
-fn sort_map_keys(keys: &Bound<PyList>, len: usize) -> Vec<(PyBackedStr, usize)> {
+fn sort_map_keys(keys: &Bound<PyList>, len: usize) -> Result<Vec<(PyBackedStr, usize)>> {
     // Returns key and index.
     let mut keys_str = Vec::with_capacity(len);
     for i in 0..len {
-        let item = keys.get_item(i).unwrap();
-        let key = item.downcast::<PyString>().unwrap().to_owned();
-        let backed_str = PyBackedStr::try_from(key).unwrap();
+        let item = keys.get_item(i)?;
+        let key = match item.downcast::<PyString>() {
+            Ok(k) => k.to_owned(),
+            Err(_) => return Err(anyhow!("Map keys must be strings")),
+        };
+        let backed_str = match PyBackedStr::try_from(key) {
+            Ok(bs) => bs,
+            Err(_) => return Err(anyhow!("Failed to convert PyString to PyBackedStr")),
+        };
         keys_str.push((backed_str, i));
     }
 
     if keys_str.len() < 2 {
-        return keys_str;
+        return Ok(keys_str);
     }
 
     keys_str.sort_by(|a, b| {
@@ -78,7 +84,7 @@ fn sort_map_keys(keys: &Bound<PyList>, len: usize) -> Vec<(PyBackedStr, usize)>
         }
     });
 
-    keys_str
+    Ok(keys_str)
 }
 
 fn get_bytes_from_py_any<'py>(obj: &'py Bound<'py, PyAny>) -> PyResult<&'py [u8]> {
@@ -252,7 +258,7 @@ fn encode_dag_cbor_from_pyobject<'py, W: Write>(
         Ok(())
     } else if let Ok(map) = obj.downcast::<PyDict>() {
         let len = map.len();
-        let keys = sort_map_keys(&map.keys(), len);
+        let keys = sort_map_keys(&map.keys(), len)?;
         let values = map.values();
 
         encode::write_u64(w, MajorKind::Map, len as u64)?;