From 2e9bfde9b41aa8805e3fec6849913f99fd3e7cd4 Mon Sep 17 00:00:00 2001
From: ryan <ryan@distrust.co>
Date: Mon, 11 Sep 2023 23:00:30 -0500
Subject: [PATCH] keyfork-seed: fixup tests, add Client

---
 keyfork-seed/src/client.rs | 21 +++++++++++++++++++++
 keyfork-seed/src/main.rs   | 15 ++++++++-------
 keyfork-seed/src/tests.rs  | 34 ++++++++++++++++++++++------------
 3 files changed, 51 insertions(+), 19 deletions(-)
 create mode 100644 keyfork-seed/src/client.rs

diff --git a/keyfork-seed/src/client.rs b/keyfork-seed/src/client.rs
new file mode 100644
index 0000000..472f738
--- /dev/null
+++ b/keyfork-seed/src/client.rs
@@ -0,0 +1,21 @@
+use keyfork_frame::*;
+use crate::Result;
+use std::os::unix::net::UnixStream;
+use keyfork_derive_util::request::*;
+
+#[derive(Debug)]
+pub struct Client {
+    socket: UnixStream,
+}
+
+impl Client {
+    pub fn new(socket: UnixStream) -> Self {
+        Self { socket }
+    }
+
+    pub fn request(&mut self, req: &DerivationRequest) -> Result<DerivationResponse> {
+        try_encode_to(&bincode::serialize(&req)?, &mut self.socket)?;
+        let resp = try_decode_from(&mut self.socket)?;
+        bincode::deserialize(&resp).map_err(From::from)
+    }
+}
diff --git a/keyfork-seed/src/main.rs b/keyfork-seed/src/main.rs
index ff79e54..442011b 100644
--- a/keyfork-seed/src/main.rs
+++ b/keyfork-seed/src/main.rs
@@ -4,6 +4,7 @@ use std::path::PathBuf;
 
 mod cli;
 mod socket;
+mod client;
 
 #[cfg(test)]
 mod tests;
@@ -35,18 +36,18 @@ pub enum Error {
     FrameDec(#[from] DecodeError),
 }
 
-fn main() -> Result<(), Error> {
+pub type Result<T, E = Error> = std::result::Result<T, E>;
+
+fn main() -> Result<()> {
     let args = cli::get_args();
-    let mut socket = socket::get_socket()?;
+    let socket = socket::get_socket()?;
+    let mut client = client::Client::new(socket);
     let path = args.get_one::<DerivationPath>("path").expect("required");
     let algo = args
         .get_one::<DerivationAlgorithm>("algorithm")
         .expect("required");
-    let req = DerivationRequest::new(algo.clone(), path.clone());
-    let ser_req = bincode::serialize(&req)?;
-    try_encode_to(&ser_req, &mut socket)?;
-    let ser_response = try_decode_from(&mut socket)?;
-    let response: DerivationResponse = bincode::deserialize(&ser_response)?;
+    let request = DerivationRequest::new(algo.clone(), path.clone());
+    let response = client.request(&request)?;
     dbg!(&response);
     Ok(())
 }
diff --git a/keyfork-seed/src/tests.rs b/keyfork-seed/src/tests.rs
index 6616100..90bdd17 100644
--- a/keyfork-seed/src/tests.rs
+++ b/keyfork-seed/src/tests.rs
@@ -1,37 +1,47 @@
+use crate::client::Client;
 use hex_literal::hex;
 use keyfork_derive_util::{request::*, DerivationPath};
-use keyfork_frame::*;
 use std::{os::unix::net::UnixStream, str::FromStr};
 use tempdir::TempDir;
 use tokio::runtime::Builder;
+use std::sync::mpsc::channel;
 
 #[test]
 fn it_works() {
     // Test literals taken from keyfork-derive-util.
+    // Setup
     let entropy = &hex!("000102030405060708090a0b0c0d0e0f")[..];
     let mnemonic = keyforkd::Mnemonic::from_entropy(entropy, Default::default()).unwrap();
     let rt = Builder::new_multi_thread().enable_io().build().unwrap();
     let tempdir = TempDir::new("keyfork-seed").unwrap();
     let socket_path = tempdir.path().join("keyforkd.sock");
+    let (tx, rx) = channel();
+
     let handle = rt.spawn({
         let socket_path = socket_path.clone();
         async move {
-            keyforkd::start_and_run_server_on(mnemonic, &socket_path)
-                .await
-                .unwrap();
+            let mut server = keyforkd::UnixServer::bind(&socket_path).unwrap();
+            // Connections can be pending for a few seconds, so signal to the main
+            // test we're ready to start accepting.
+            tx.send(()).unwrap();
+            let service = keyforkd::ServiceBuilder::new()
+                .layer(keyforkd::middleware::BincodeLayer::new())
+                .service(keyforkd::Keyforkd::new(mnemonic));
+            server.run(service).await.unwrap();
         }
     });
-    // TODO: send a channel to start_and_run_server_on, or extract logic out of server,
-    // so we can send a "started" message back to main thread.
-    std::thread::sleep(std::time::Duration::from_secs(2));
-    let mut socket = UnixStream::connect(&socket_path).unwrap();
+
+    rx.recv().unwrap();
+    let socket = UnixStream::connect(&socket_path).unwrap();
+    let mut client = Client::new(socket);
     let req = DerivationRequest::new(
         DerivationAlgorithm::Ed25519,
         DerivationPath::from_str("m/0'/1'/2'/2'/1000000000'").unwrap(),
     );
-    try_encode_to(&bincode::serialize(&req).unwrap(), &mut socket).unwrap();
-    let response: DerivationResponse =
-        bincode::deserialize(&try_decode_from(&mut socket).unwrap()).unwrap();
-    assert_eq!(response.data, hex!("8f94d394a8e8fd6b1bc2f3f49f5c47e385281d5c17e65324b0f62483e37e8793"));
+    let response = client.request(&req).unwrap();
+    assert_eq!(
+        response.data,
+        hex!("8f94d394a8e8fd6b1bc2f3f49f5c47e385281d5c17e65324b0f62483e37e8793")
+    );
     handle.abort();
 }