1 files changed, 41 insertions, 25 deletions

diff --git a/src/util.rs b/src/util.rs
index a5dd1e5..26942cf 100644
--- a/src/util.rs
+++ b/src/util.rs
@@ -5,8 +5,7 @@ use std::ffi::{CStr, CString};
 use std::os::raw::{c_char, c_int};
 
 use libc::{c_void, free};
-use rand_core::RngCore;
-use rand_os::OsRng;
+use rand_core::{OsRng, RngCore};
 
 use crate::error::{Error, LibraryError};
 
@@ -31,26 +30,38 @@ pub enum LogLevel {
     DebugL2,
 }
 
+pub fn str_from_ptr<'a>(ptr: *const c_char) -> Result<&'a str, Error> {
+    unsafe { CStr::from_ptr(ptr) }.to_str().map_err(Error::from)
+}
+
 pub fn owned_str_from_ptr(ptr: *const c_char) -> Result<String, Error> {
-    unsafe { CStr::from_ptr(ptr) }
-        .to_str()
-        .map(String::from)
-        .map_err(Error::from)
+    str_from_ptr(ptr).map(ToOwned::to_owned)
 }
 
-pub fn result_from_string(ptr: *const c_char) -> Result<String, Error> {
+pub fn run_with_string<R, F>(ptr: *const c_char, op: F) -> Result<R, Error>
+where
+    F: FnOnce(&str) -> Result<R, Error>,
+{
     if ptr.is_null() {
-        return Err(Error::UnexpectedError);
+        return Err(Error::UnexpectedError(
+            "libnitrokey returned a null pointer".to_owned(),
+        ));
     }
-    let s = owned_str_from_ptr(ptr)?;
+    let result = str_from_ptr(ptr).and_then(op);
     unsafe { free(ptr as *mut c_void) };
-    // An empty string can both indicate an error or be a valid return value.  In this case, we
-    // have to check the last command status to decide what to return.
-    if s.is_empty() {
-        get_last_result().map(|_| s)
-    } else {
-        Ok(s)
-    }
+    result
+}
+
+pub fn result_from_string(ptr: *const c_char) -> Result<String, Error> {
+    run_with_string(ptr, |s| {
+        // An empty string can both indicate an error or be a valid return value.  In this case, we
+        // have to check the last command status to decide what to return.
+        if s.is_empty() {
+            get_last_result().map(|_| s.to_owned())
+        } else {
+            Ok(s.to_owned())
+        }
+    })
 }
 
 pub fn result_or_error<T>(value: T) -> Result<T, Error> {
@@ -70,20 +81,25 @@ pub fn get_last_result() -> Result<(), Error> {
 }
 
 pub fn get_last_error() -> Error {
-    match get_last_result() {
-        Ok(()) => Error::UnexpectedError,
-        Err(err) => err,
-    }
+    get_last_result().err().unwrap_or_else(|| {
+        Error::UnexpectedError("Expected an error, but command status is zero".to_owned())
+    })
 }
 
-pub fn generate_password(length: usize) -> Result<Vec<u8>, Error> {
-    let mut data = vec![0u8; length];
-    OsRng.fill_bytes(&mut data[..]);
-    Ok(data)
+pub fn generate_password(length: usize) -> Result<CString, Error> {
+    loop {
+        // Randomly generate a password until we get a string *without* null bytes.  Otherwise
+        // the string would be cut off prematurely due to null-termination in C.
+        let mut data = vec![0u8; length];
+        OsRng.fill_bytes(&mut data[..]);
+        if let Ok(s) = CString::new(data) {
+            return Ok(s);
+        }
+    }
 }
 
 pub fn get_cstring<T: Into<Vec<u8>>>(s: T) -> Result<CString, Error> {
-    CString::new(s).or_else(|_| Err(LibraryError::InvalidString.into()))
+    CString::new(s).map_err(|_| LibraryError::InvalidString.into())
 }
 
 impl Into<i32> for LogLevel {