Add more session cache support

3 files changed, 148 insertions, 3 deletions

diff --git a/openssl/src/ssl/callbacks.rs b/openssl/src/ssl/callbacks.rs
index 9a3d3de1..0a38952b 100644
--- a/openssl/src/ssl/callbacks.rs
+++ b/openssl/src/ssl/callbacks.rs
@@ -5,6 +5,8 @@ use std::ptr;
 use std::slice;
 use std::mem;
 use foreign_types::ForeignTypeRef;
+#[cfg(any(all(feature = "v110", ossl110), all(feature = "v111", ossl111)))]
+use foreign_types::ForeignType;
 
 use error::ErrorStack;
 use dh::Dh;
@@ -15,6 +17,8 @@ use ssl::{get_callback_idx, get_ssl_callback_idx, SniError, SslAlert, SslRef};
 #[cfg(any(all(feature = "v102", ossl102), all(feature = "v110", ossl110),
           all(feature = "v111", ossl111)))]
 use ssl::AlpnError;
+#[cfg(any(all(feature = "v110", ossl110), all(feature = "v111", ossl111)))]
+use ssl::SslSession;
 use x509::X509StoreContextRef;
 
 pub extern "C" fn raw_verify<F>(preverify_ok: c_int, x509_ctx: *mut ffi::X509_STORE_CTX) -> c_int
@@ -274,3 +278,24 @@ where
         }
     }
 }
+
+#[cfg(any(all(feature = "v110", ossl110), all(feature = "v111", ossl111)))]
+pub unsafe extern "C" fn raw_new_session<F>(
+    ssl: *mut ffi::SSL,
+    session: *mut ffi::SSL_SESSION,
+) -> c_int
+where
+    F: Fn(&mut SslRef, SslSession) + 'static + Sync + Send,
+{
+    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 callback = &*(callback as *mut F);
+
+    let ssl = SslRef::from_ptr_mut(ssl);
+    let session = SslSession::from_ptr(session);
+
+    callback(ssl, session);
+
+    // the return code doesn't indicate error vs success, but whether or not we consumed the session
+    1
+}
diff --git a/openssl/src/ssl/mod.rs b/openssl/src/ssl/mod.rs
index 8e483015..7308535b 100644
--- a/openssl/src/ssl/mod.rs
+++ b/openssl/src/ssl/mod.rs
@@ -305,19 +305,55 @@ bitflags! {
         /// Verifies that the peer's certificate is trusted.
         ///
         /// On the server side, this will cause OpenSSL to request a certificate from the client.
-        const PEER = ::ffi::SSL_VERIFY_PEER;
+        const PEER = ffi::SSL_VERIFY_PEER;
 
         /// Disables verification of the peer's certificate.
         ///
         /// On the server side, this will cause OpenSSL to not request a certificate from the
         /// client. On the client side, the certificate will be checked for validity, but the
         /// negotiation will continue regardless of the result of that check.
-        const NONE = ::ffi::SSL_VERIFY_NONE;
+        const NONE = ffi::SSL_VERIFY_NONE;
 
         /// On the server side, abort the handshake if the client did not send a certificate.
         ///
         /// This should be paired with `SSL_VERIFY_PEER`. It has no effect on the client side.
-        const FAIL_IF_NO_PEER_CERT = ::ffi::SSL_VERIFY_FAIL_IF_NO_PEER_CERT;
+        const FAIL_IF_NO_PEER_CERT = ffi::SSL_VERIFY_FAIL_IF_NO_PEER_CERT;
+    }
+}
+
+bitflags! {
+    /// Options controlling the behavior of session caching.
+    pub struct SslSessionCacheMode: c_long {
+        /// No session caching for the client or server takes place.
+        const OFF = ffi::SSL_SESS_CACHE_OFF;
+
+        /// Enable session caching on the client side.
+        ///
+        /// OpenSSL has no way of identifying the proper session to reuse automatically, so the
+        /// application is responsible for setting it explicitly via [`SslRef::set_session`].
+        ///
+        /// [`SslRef::set_session`]: struct.SslRef.html#method.set_session
+        const CLIENT = ffi::SSL_SESS_CACHE_CLIENT;
+
+        /// Enable session caching on the server side.
+        ///
+        /// This is the default mode.
+        const SERVER = ffi::SSL_SESS_CACHE_SERVER;
+
+        /// Enable session caching on both the client and server side.
+        const BOTH = ffi::SSL_SESS_CACHE_BOTH;
+
+        /// Disable automatic removal of expired sessions from the session cache.
+        const NO_AUTO_CLEAR = ffi::SSL_SESS_CACHE_NO_AUTO_CLEAR;
+
+        /// Disable use of the internal session cache for session lookups.
+        const NO_INTERNAL_LOOKUP = ffi::SSL_SESS_CACHE_NO_INTERNAL_LOOKUP;
+
+        /// Disable use of the internal session cache for session storage.
+        const NO_INTERNAL_STORE = ffi::SSL_SESS_CACHE_NO_INTERNAL_STORE;
+
+        /// Disable use of the internal session cache for storage and lookup.
+        const NO_INTERNAL = ffi::SSL_SESS_CACHE_NO_INTERNAL;
     }
 }
 
@@ -1128,6 +1164,54 @@ impl SslContextBuilder {
         }
     }
 
+    /// Sets the callback which is called when new sessions are negotiated.
+    ///
+    /// This is used by clients to implement session caching. While for TLSv1.2 the session is
+    /// available to access via [`SslRef::session`] immediately after the handshake completes, this
+    /// is not the case for TLSv1.3. There, a session is not generally available immediately, and
+    /// the server may provide multiple session tokens to the client over a single session. The new
+    /// session callback is a portable way to deal with both cases.
+    ///
+    /// Note that session caching must be enabled for the callback to be invoked, and it defaults
+    /// off for clients. [`set_session_cache_mode`] controls that behavior.
+    ///
+    /// This corresponds to [`SSL_CTX_sess_set_new_cb`].
+    ///
+    /// Requires OpenSSL 1.1.0 or 1.1.1 and the corresponding Cargo feature.
+    ///
+    /// [`SslRef::session`]: struct.SslRef.html#method.session
+    /// [`set_session_cache_mode`]: #method.set_session_cache_mode
+    /// [`SSL_CTX_sess_set_new_cb`]: https://www.openssl.org/docs/manmaster/man3/SSL_CTX_sess_set_new_cb.html
+    #[cfg(any(all(feature = "v110", ossl110), all(feature = "v111", ossl111)))]
+    pub fn set_new_session_callback<F>(&mut self, callback: F)
+    where
+        F: Fn(&mut SslRef, SslSession) + 'static + Sync + Send,
+    {
+        unsafe {
+            let callback = Box::new(callback);
+            ffi::SSL_CTX_set_ex_data(
+                self.as_ptr(),
+                get_callback_idx::<F>(),
+                Box::into_raw(callback) as *mut _,
+            );
+            ffi::SSL_CTX_sess_set_new_cb(self.as_ptr(), Some(callbacks::raw_new_session::<F>));
+        }
+    }
+
+    /// Sets the session caching mode use for connections made with the context.
+    ///
+    /// Returns the previous session caching mode.
+    ///
+    /// This corresponds to [`SSL_CTX_set_session_cache_mode`].
+    ///
+    /// [`SSL_CTX_set_session_cache_mode`]: https://www.openssl.org/docs/man1.1.0/ssl/SSL_CTX_get_session_cache_mode.html
+    pub fn set_session_cache_mode(&mut self, mode: SslSessionCacheMode) -> SslSessionCacheMode {
+        unsafe {
+            let bits = ffi::SSL_CTX_set_session_cache_mode(self.as_ptr(), mode.bits());
+            SslSessionCacheMode { bits }
+        }
+    }
+
     /// Sets the extra data at the specified index.
     ///
     /// This can be used to provide data to callbacks registered with the context. Use the
diff --git a/openssl/src/ssl/test.rs b/openssl/src/ssl/test.rs
index 51ae6cae..17bc37d9 100644
--- a/openssl/src/ssl/test.rs
+++ b/openssl/src/ssl/test.rs
@@ -20,6 +20,8 @@ use ocsp::{OcspResponse, OcspResponseStatus};
 use ssl;
 use ssl::{Error, HandshakeError, ShutdownResult, Ssl, SslAcceptor, SslConnector, SslContext,
           SslFiletype, SslMethod, SslStream, SslVerifyMode, StatusType};
+#[cfg(any(all(feature = "v110", ossl110), all(feature = "v111", ossl111)))]
+use ssl::SslSessionCacheMode;
 use x509::{X509, X509Name, X509StoreContext, X509VerifyResult};
 #[cfg(any(all(feature = "v102", ossl102), all(feature = "v110", ossl110),
           all(feature = "v111", ossl111)))]
@@ -1245,6 +1247,40 @@ fn status_callbacks() {
     guard.join().unwrap();
 }
 
+#[test]
+#[cfg(any(all(feature = "v110", ossl110), all(feature = "v111", ossl111)))]
+fn new_session_callback() {
+    static CALLED_BACK: AtomicBool = ATOMIC_BOOL_INIT;
+
+    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
+    let port = listener.local_addr().unwrap().port();
+
+    thread::spawn(move || {
+        let stream = listener.accept().unwrap().0;
+        let mut ctx = SslContext::builder(SslMethod::tls()).unwrap();
+        ctx.set_certificate_file(&Path::new("test/cert.pem"), SslFiletype::PEM)
+            .unwrap();
+        ctx.set_private_key_file(&Path::new("test/key.pem"), SslFiletype::PEM)
+            .unwrap();
+        ctx.set_session_id_context(b"foo").unwrap();
+        let ssl = Ssl::new(&ctx.build()).unwrap();
+        let mut stream = ssl.accept(stream).unwrap();
+        stream.write_all(&[0]).unwrap();
+    });
+
+    let stream = TcpStream::connect(("127.0.0.1", port)).unwrap();
+    let mut ctx = SslContext::builder(SslMethod::tls()).unwrap();
+    ctx.set_session_cache_mode(SslSessionCacheMode::CLIENT | SslSessionCacheMode::NO_INTERNAL);
+    ctx.set_new_session_callback(|_, _| CALLED_BACK.store(true, Ordering::SeqCst));
+    let ssl = Ssl::new(&ctx.build()).unwrap();
+    let mut stream = ssl.connect(stream).unwrap();
+    // read 1 byte to make sure the session is received for TLSv1.3
+    let mut buf = [0];
+    stream.read_exact(&mut buf).unwrap();
+
+    assert!(CALLED_BACK.load(Ordering::SeqCst));
+}
+
 fn _check_kinds() {
     fn is_send<T: Send>() {}
     fn is_sync<T: Sync>() {}