Fix generic connections (connect()) to return correct device

This patch fixes the generic connect() method to return a DeviceWrapper
of the correct type.  This is enabled by the NK_get_device_model()
method introduced in libnitrokey v3.4.

3 files changed, 26 insertions, 3 deletions

diff --git a/TODO.md b/TODO.md
index 58d923b..38044ae 100644
--- a/TODO.md
+++ b/TODO.md
@@ -36,7 +36,6 @@
 - Find a nicer syntax for the `write_config` test.
 - Prevent construction of internal types.
 - Add Storage-only examples to the `DeviceWrapper` documentation.
-- Fix generic connection (`get_connected_device`).
 - More specific error checking in the tests.
 - Differentiate empty strings and errors (see `result_from_string`).
 - Check integer conversions.
diff --git a/src/device.rs b/src/device.rs
index 326df5d..4f96860 100644
--- a/src/device.rs
+++ b/src/device.rs
@@ -445,9 +445,25 @@ pub fn connect() -> Result<DeviceWrapper, CommandError> {
     }
 }
 
+fn get_connected_model() -> Option<Model> {
+    unsafe {
+        match nitrokey_sys::NK_get_device_model() {
+            nitrokey_sys::NK_device_model_NK_PRO => Some(Model::Pro),
+            nitrokey_sys::NK_device_model_NK_STORAGE => Some(Model::Storage),
+            _ => None,
+        }
+    }
+}
+
+fn create_device_wrapper(model: Model) -> DeviceWrapper {
+    match model {
+        Model::Pro => DeviceWrapper::Pro(Pro {}),
+        Model::Storage => DeviceWrapper::Storage(Storage {}),
+    }
+}
+
 fn get_connected_device() -> Option<DeviceWrapper> {
-    // TODO: check connected device
-    Some(DeviceWrapper::Pro(Pro {}))
+    get_connected_model().map(create_device_wrapper)
 }
 
 fn connect_model(model: Model) -> bool {
diff --git a/src/tests/device.rs b/src/tests/device.rs
index 60ca33a..c2c5336 100644
--- a/src/tests/device.rs
+++ b/src/tests/device.rs
@@ -33,6 +33,10 @@ fn connect_pro() {
     assert!(::connect().is_ok());
     assert!(::Pro::connect().is_ok());
     assert!(::Storage::connect().is_err());
+    match ::connect().unwrap() {
+        ::DeviceWrapper::Pro(_) => assert!(true),
+        ::DeviceWrapper::Storage(_) => assert!(false),
+    };
 }
 
 #[test]
@@ -41,6 +45,10 @@ fn connect_storage() {
     assert!(::connect().is_ok());
     assert!(::Pro::connect().is_err());
     assert!(::Storage::connect().is_ok());
+    match ::connect().unwrap() {
+        ::DeviceWrapper::Pro(_) => assert!(false),
+        ::DeviceWrapper::Storage(_) => assert!(true),
+    };
 }
 
 fn assert_empty_serial_number() {