Make some changes for review comments

1 files changed, 24 insertions, 21 deletions

diff --git a/openssl/src/ssl/mod.rs b/openssl/src/ssl/mod.rs
index b58bbf41..0a380b6f 100644
--- a/openssl/src/ssl/mod.rs
+++ b/openssl/src/ssl/mod.rs
@@ -76,7 +76,7 @@ use libc::{c_int, c_void, c_long, c_ulong};
 use libc::{c_char, c_uchar, c_uint};
 use std::any::Any;
 use std::any::TypeId;
-use std::borrow::{Borrow, Cow};
+use std::borrow::Borrow;
 use std::cmp;
 use std::collections::HashMap;
 use std::ffi::{CStr, CString};
@@ -307,35 +307,32 @@ extern "C" fn raw_verify<F>(preverify_ok: c_int, x509_ctx: *mut ffi::X509_STORE_
     }
 }
 
+#[cfg(not(osslconf = "OPENSSL_NO_PSK"))]
 extern "C" fn raw_psk<F>(ssl: *mut ffi::SSL,
                          hint: *const c_char,
                          identity: *mut c_char,
-                         _max_identity_len: c_uint,
+                         max_identity_len: c_uint,
                          psk: *mut c_uchar,
-                         _max_psk_len: c_uint) -> c_uint
-    where F: Fn(&mut SslRef, &str) -> Result<(String, Vec<u8>), ErrorStack> + Any + 'static + Sync + Send
+                         max_psk_len: c_uint) -> c_uint
+    where F: Fn(&mut SslRef, Option<&[u8]>, &mut [u8], &mut [u8]) -> Result<usize, ErrorStack> + Any + 'static + Sync + Send
 {
     unsafe {
         let ssl_ctx = ffi::SSL_get_SSL_CTX(ssl as *const _);
         let callback = ffi::SSL_CTX_get_ex_data(ssl_ctx, get_callback_idx::<F>());
         let ssl = SslRef::from_ptr_mut(ssl);
         let callback = &*(callback as *mut F);
-        let cstr_hint = if hint != ptr::null() { Cow::Borrowed(CStr::from_ptr(hint)) } else { Cow::Owned(CString::new("").unwrap()) };
-        if let Ok(s) = cstr_hint.to_str() {
-            match callback(ssl, s) {
-                Ok((identity_out, psk_out)) => {
-                    if let Ok(id) = CString::new(identity_out) {
-                        //TODO: validate max_identity_len, max_psk_len
-                        let id = id.into_bytes_with_nul();
-                        ptr::copy(id.as_ptr() as *mut i8, identity, id.len());
-                        ptr::copy(psk_out.as_ptr(), psk, psk_out.len());
-                        psk_out.len() as u32
-                    } else { 0 }
-                }
-                Err(_) => 0,
-            }
+        let hint = if hint != ptr::null() {
+            Some(CStr::from_ptr(hint).to_bytes())
         } else {
-            0
+            None
+        };
+        // Give the callback mutable slices into which it can write the identity and psk.
+        let identity_sl = slice::from_raw_parts_mut(identity as *mut u8,
+                                                    max_identity_len as usize);
+        let psk_sl = slice::from_raw_parts_mut(psk as *mut u8, max_psk_len as usize);
+        match callback(ssl, hint, identity_sl, psk_sl) {
+            Ok(psk_len) => psk_len as u32,
+            _ => 0,
         }
     }
 }
@@ -1010,15 +1007,21 @@ impl SslContextBuilder {
         }
     }
 
+    /// Sets the callback for providing an identity and pre-shared key for a TLS-PSK client.
+    ///
+    /// The callback will be called with the SSL context, an identity hint if one was provided
+    /// by the server, a mut slice for each of the identity and pre-shared key bytes. The identity
+    /// must be written as a null-terminated C string.
+    #[cfg(not(osslconf = "OPENSSL_NO_PSK"))]
     pub fn set_psk_callback<F>(&mut self, callback: F)
-        where F: Fn(&mut SslRef, &str) -> Result<(String, Vec<u8>), ErrorStack> + Any + 'static + Sync + Send
+        where F: Fn(&mut SslRef, Option<&[u8]>, &mut [u8], &mut [u8]) -> Result<usize, ErrorStack> + Any + 'static + Sync + Send
     {
         unsafe {
             let callback = Box::new(callback);
             ffi::SSL_CTX_set_ex_data(self.as_ptr(),
                                      get_callback_idx::<F>(),
                                      mem::transmute(callback));
-            ffi::SSL_CTX_set_psk_client_callback(self.as_ptr(), raw_psk::<F>)
+            ffi::SSL_CTX_set_psk_client_callback(self.as_ptr(), Some(raw_psk::<F>))
         }
     }