From f60b77254ab18bf9cc807ea620c31af79a0c8c4f Mon Sep 17 00:00:00 2001
From: ryan <ryan@distrust.co>
Date: Mon, 11 Sep 2023 19:44:47 -0500
Subject: [PATCH] keyforkd: add required 2 path segments

---
 keyforkd/src/middleware.rs |  5 ++-
 keyforkd/src/service.rs    | 66 ++++++++++++++++++++++++++++++++++++--
 2 files changed, 66 insertions(+), 5 deletions(-)

diff --git a/keyforkd/src/middleware.rs b/keyforkd/src/middleware.rs
index 7f2f7ea..b64726f 100644
--- a/keyforkd/src/middleware.rs
+++ b/keyforkd/src/middleware.rs
@@ -103,7 +103,7 @@ mod tests {
     impl Test {
         fn new() -> Self {
             Self {
-                field: "hello world!".to_string()
+                field: "hello world!".to_string(),
             }
         }
     }
@@ -135,8 +135,7 @@ mod tests {
 
     #[tokio::test]
     async fn can_serde_responses() {
-        let content = serialize(&Test::new())
-        .unwrap();
+        let content = serialize(&Test::new()).unwrap();
         let mut service = ServiceBuilder::new()
             .layer(BincodeLayer::<Test>::new())
             .service(App);
diff --git a/keyforkd/src/service.rs b/keyforkd/src/service.rs
index 9362338..3d62ea8 100644
--- a/keyforkd/src/service.rs
+++ b/keyforkd/src/service.rs
@@ -7,6 +7,15 @@ use tower::Service;
 // NOTE: All values implemented in Keyforkd must implement Clone with low overhead, either by
 // using an Arc or by having a small signature. This is because Service<T> takes &mut self.
 
+#[derive(thiserror::Error, Debug)]
+pub enum KeyforkdRequestError {
+    #[error("Invalid derivation length: Expected: 2, actual: {0}")]
+    InvalidDerivationLength(usize),
+
+    #[error("Derivation error: {0}")]
+    Derivation(#[from] DerivationError),
+}
+
 #[derive(Clone, Debug)]
 pub struct Keyforkd {
     mnemonic: Arc<Mnemonic>,
@@ -23,7 +32,7 @@ impl Keyforkd {
 impl Service<DerivationRequest> for Keyforkd {
     type Response = DerivationResponse;
 
-    type Error = DerivationError;
+    type Error = KeyforkdRequestError;
 
     type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
 
@@ -37,7 +46,14 @@ impl Service<DerivationRequest> for Keyforkd {
     #[cfg_attr(feature = "tracing", tracing::instrument(skip(self)))]
     fn call(&mut self, req: DerivationRequest) -> Self::Future {
         let mnemonic = self.mnemonic.clone();
-        Box::pin(async { req.derive_with_mnemonic(&mnemonic) })
+        Box::pin(async {
+            let len = req.path().len();
+            if len < 2 {
+                return Err(KeyforkdRequestError::InvalidDerivationLength(len));
+            }
+            req.derive_with_mnemonic(&mnemonic)
+                .map_err(KeyforkdRequestError::from)
+        })
     }
 }
 
@@ -54,6 +70,9 @@ mod tests {
     async fn properly_derives_data() {
         // Pulled from keyfork-derive-util's tests, which is more extensively tested.
         let tests = [
+            /*
+             * Note: Tests excluded because the derivation path is not deep enough
+             * for the API's preferences.
             (
                 &hex!("000102030405060708090a0b0c0d0e0f")[..],
                 DerivationPath::from_str("m").unwrap(),
@@ -68,6 +87,7 @@ mod tests {
                 hex!("68e0fe46dfb67e368c75379acec591dad19df3cde26e63b93a8e704f1dade7a3"),
                 hex!("008c8a13df77a28f3445213a0f432fde644acaa215fc72dcdf300d5efaa85d350c"),
             ),
+            */
             (
                 &hex!("000102030405060708090a0b0c0d0e0f")[..],
                 DerivationPath::from_str("m/0'/1'/2'/2'/1000000000'").unwrap(),
@@ -86,4 +106,46 @@ mod tests {
             assert_eq!(response.data, private_key)
         }
     }
+
+    #[should_panic]
+    #[tokio::test]
+    async fn errors_on_no_path() {
+        let tests = [(
+            &hex!("000102030405060708090a0b0c0d0e0f")[..],
+            DerivationPath::from_str("m").unwrap(),
+            hex!("90046a93de5380a72b5e45010748567d5ea02bbf6522f979e05c0d8d8ca9fffb"),
+            hex!("2b4be7f19ee27bbf30c667b642d5f4aa69fd169872f8fc3059c08ebae2eb19e7"),
+            hex!("00a4b2856bfec510abab89753fac1ac0e1112364e7d250545963f135f2a33188ed"),
+        )];
+        let wordlist = Wordlist::default().arc();
+        for (seed, path, _, private_key, _) in tests {
+            let mnemonic = Mnemonic::from_entropy(&seed[..], wordlist.clone()).unwrap();
+            assert_eq!(mnemonic.seed(), seed);
+            let req = DerivationRequest::new(DerivationAlgorithm::Ed25519, path);
+            let mut keyforkd = Keyforkd::new(mnemonic);
+            let response = keyforkd.ready().await.unwrap().call(req).await.unwrap();
+            assert_eq!(response.data, private_key)
+        }
+    }
+
+    #[should_panic]
+    #[tokio::test]
+    async fn errors_on_short_path() {
+        let tests = [(
+            &hex!("000102030405060708090a0b0c0d0e0f")[..],
+            DerivationPath::from_str("m/0'").unwrap(),
+            hex!("8b59aa11380b624e81507a27fedda59fea6d0b779a778918a2fd3590e16e9c69"),
+            hex!("68e0fe46dfb67e368c75379acec591dad19df3cde26e63b93a8e704f1dade7a3"),
+            hex!("008c8a13df77a28f3445213a0f432fde644acaa215fc72dcdf300d5efaa85d350c"),
+        )];
+        let wordlist = Wordlist::default().arc();
+        for (seed, path, _, private_key, _) in tests {
+            let mnemonic = Mnemonic::from_entropy(&seed[..], wordlist.clone()).unwrap();
+            assert_eq!(mnemonic.seed(), seed);
+            let req = DerivationRequest::new(DerivationAlgorithm::Ed25519, path);
+            let mut keyforkd = Keyforkd::new(mnemonic);
+            let response = keyforkd.ready().await.unwrap().call(req).await.unwrap();
+            assert_eq!(response.data, private_key)
+        }
+    }
 }