From 47a68e2929277889c466288333385bd367238a42 Mon Sep 17 00:00:00 2001
From: Benjamin Cheng <ben@bcheng.cf>
Date: Wed, 16 May 2018 17:49:36 -0400
Subject: [PATCH] Add wrapper for SSL_CTX_set_psk_server_callback

---
 openssl-sys/src/lib.rs       |  8 ++++++++
 openssl/src/ssl/callbacks.rs | 36 +++++++++++++++++++++++++++++++++++-
 openssl/src/ssl/mod.rs       | 27 +++++++++++++++++++++++++--
 3 files changed, 68 insertions(+), 3 deletions(-)

diff --git a/openssl-sys/src/lib.rs b/openssl-sys/src/lib.rs
index 61e087d0..9f970701 100644
--- a/openssl-sys/src/lib.rs
+++ b/openssl-sys/src/lib.rs
@@ -2576,6 +2576,14 @@ extern "C" {
                 -> c_uint,
         >,
     );
+    #[cfg(not(osslconf = "OPENSSL_NO_PSK"))]
+    pub fn SSL_CTX_set_psk_server_callback(
+        ssl: *mut SSL_CTX,
+        psk_server_cb: Option<
+            extern "C" fn(*mut SSL, *const c_char, *mut c_uchar, c_uint)
+                -> c_uint,
+        >,
+    );
 
     pub fn SSL_select_next_proto(
         out: *mut *mut c_uchar,
diff --git a/openssl/src/ssl/callbacks.rs b/openssl/src/ssl/callbacks.rs
index fd5b7ef5..887d450c 100644
--- a/openssl/src/ssl/callbacks.rs
+++ b/openssl/src/ssl/callbacks.rs
@@ -48,7 +48,7 @@ where
 }
 
 #[cfg(not(osslconf = "OPENSSL_NO_PSK"))]
-pub extern "C" fn raw_psk<F>(
+pub extern "C" fn raw_client_psk<F>(
     ssl: *mut ffi::SSL,
     hint: *const c_char,
     identity: *mut c_char,
@@ -84,6 +84,40 @@ where
     }
 }
 
+#[cfg(not(osslconf = "OPENSSL_NO_PSK"))]
+pub extern "C" fn raw_server_psk<F>(
+    ssl: *mut ffi::SSL,
+    identity: *const c_char,
+    psk: *mut c_uchar,
+    max_psk_len: c_uint,
+) -> c_uint
+where
+    F: Fn(&mut SslRef, Option<&[u8]>, &mut [u8]) -> Result<usize, ErrorStack>
+        + 'static
+        + Sync
+        + Send,
+{
+    unsafe {
+        let ssl = SslRef::from_ptr_mut(ssl);
+        let callback_idx = SslContext::cached_ex_index::<F>();
+
+        let callback = ssl.ssl_context()
+            .ex_data(callback_idx)
+            .expect("BUG: psk callback missing") as *const F;
+        let identity = if identity != ptr::null() {
+            Some(CStr::from_ptr(identity).to_bytes())
+        } else {
+            None
+        };
+        // Give the callback mutable slices into which it can write the psk.
+        let psk_sl = slice::from_raw_parts_mut(psk as *mut u8, max_psk_len as usize);
+        match (*callback)(ssl, identity, psk_sl) {
+            Ok(psk_len) => psk_len as u32,
+            _ => 0,
+        }
+    }
+}
+
 pub extern "C" fn ssl_raw_verify<F>(
     preverify_ok: c_int,
     x509_ctx: *mut ffi::X509_STORE_CTX,
diff --git a/openssl/src/ssl/mod.rs b/openssl/src/ssl/mod.rs
index 8dc605ed..37f5086c 100644
--- a/openssl/src/ssl/mod.rs
+++ b/openssl/src/ssl/mod.rs
@@ -1226,7 +1226,7 @@ impl SslContextBuilder {
     ///
     /// [`SSL_CTX_set_psk_client_callback`]: https://www.openssl.org/docs/man1.0.2/ssl/SSL_CTX_set_psk_client_callback.html
     #[cfg(not(osslconf = "OPENSSL_NO_PSK"))]
-    pub fn set_psk_callback<F>(&mut self, callback: F)
+    pub fn set_psk_client_callback<F>(&mut self, callback: F)
     where
         F: Fn(&mut SslRef, Option<&[u8]>, &mut [u8], &mut [u8]) -> Result<usize, ErrorStack>
             + 'static
@@ -1235,7 +1235,30 @@ impl SslContextBuilder {
     {
         unsafe {
             self.set_ex_data(SslContext::cached_ex_index::<F>(), callback);
-            ffi::SSL_CTX_set_psk_client_callback(self.as_ptr(), Some(raw_psk::<F>));
+            ffi::SSL_CTX_set_psk_client_callback(self.as_ptr(), Some(raw_client_psk::<F>));
+        }
+    }
+
+    /// Sets the callback for providing an identity and pre-shared key for a TLS-PSK server.
+    ///
+    /// The callback will be called with the SSL context, an identity provided by the client,
+    /// and, a mutable slice for the pre-shared key bytes. The callback returns the number of
+    /// bytes in the pre-shared key.
+    ///
+    /// This corresponds to [`SSL_CTX_set_psk_server_callback`].
+    ///
+    /// [`SSL_CTX_set_psk_server_callback`]: https://www.openssl.org/docs/man1.0.2/ssl/SSL_CTX_set_psk_server_callback.html
+    #[cfg(not(osslconf = "OPENSSL_NO_PSK"))]
+    pub fn set_psk_server_callback<F>(&mut self, callback: F)
+    where
+        F: Fn(&mut SslRef, Option<&[u8]>, &mut [u8]) -> Result<usize, ErrorStack>
+            + 'static
+            + Sync
+            + Send,
+    {
+        unsafe {
+            self.set_ex_data(SslContext::cached_ex_index::<F>(), callback);
+            ffi::SSL_CTX_set_psk_server_callback(self.as_ptr(), Some(raw_server_psk::<F>));
         }
     }