From 53b2f6615b622a061d18a34024944a6125e3f3c0 Mon Sep 17 00:00:00 2001
From: Pedro Camboim <72642597+pedrxlz@users.noreply.github.com>
Date: Thu, 6 Feb 2025 16:14:34 -0300
Subject: [PATCH] feat: added support for 64 bytes pvk (#97)

---
 packages/kos/src/chains/klv/mod.rs | 30 ++++++++++++++++++++++++++++++
 packages/kos/src/chains/util.rs    |  5 ++++-
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/packages/kos/src/chains/klv/mod.rs b/packages/kos/src/chains/klv/mod.rs
index 034d87f..54f4635 100644
--- a/packages/kos/src/chains/klv/mod.rs
+++ b/packages/kos/src/chains/klv/mod.rs
@@ -189,6 +189,36 @@ mod test {
         );
     }
 
+    #[test]
+    fn test_pvk_32() {
+        let pvk = hex::decode("8734062c1158f26a3ca8a4a0da87b527a7c168653f7f4c77045e5cf571497d9d")
+            .unwrap();
+
+        let pbk = crate::chains::klv::KLV {}.get_pbk(pvk).unwrap();
+
+        let address = crate::chains::klv::KLV {}.get_address(pbk.clone()).unwrap();
+        assert_eq!(
+            address,
+            "klv1usdnywjhrlv4tcyu6stxpl6yvhplg35nepljlt4y5r7yppe8er4qujlazy"
+        );
+        assert_eq!(pbk.len(), 32);
+    }
+
+    #[test]
+    fn test_pvk_64() {
+        let pvk = hex::decode("8734062c1158f26a3ca8a4a0da87b527a7c168653f7f4c77045e5cf571497d9de41b323a571fd955e09cd41660ff4465c3f44693c87f2faea4a0fc408727c8ea")
+            .unwrap();
+
+        let pbk = crate::chains::klv::KLV {}.get_pbk(pvk).unwrap();
+
+        let address = crate::chains::klv::KLV {}.get_address(pbk.clone()).unwrap();
+        assert_eq!(
+            address,
+            "klv1usdnywjhrlv4tcyu6stxpl6yvhplg35nepljlt4y5r7yppe8er4qujlazy"
+        );
+        assert_eq!(pbk.len(), 32);
+    }
+
     #[test]
     fn test_sign_raw() {
         let mnemonic =
diff --git a/packages/kos/src/chains/util.rs b/packages/kos/src/chains/util.rs
index 2ed63ab..ea5b991 100644
--- a/packages/kos/src/chains/util.rs
+++ b/packages/kos/src/chains/util.rs
@@ -11,5 +11,8 @@ pub fn slice_from_vec<const N: usize>(vec: &[u8]) -> Result<[u8; N], ChainError>
 }
 
 pub fn private_key_from_vec<const N: usize>(vec: &[u8]) -> Result<[u8; N], ChainError> {
-    slice_from_vec::<N>(vec).map_err(|_| ChainError::InvalidPrivateKey)
+    // If input is longer than N, take first N bytes
+    let slice = if vec.len() > N { &vec[..N] } else { vec };
+
+    slice_from_vec::<N>(slice).map_err(|_| ChainError::InvalidPrivateKey)
 }