Add PSK support

Co-authored-by: Natnatee Dokmai <[email protected]>

Author: dongcarl

mbedtls/Cargo.toml

@@ -85,6 +85,10 @@ required-features = ["std"]
 name = "ssl_conf_ca_cb"
 required-features = ["std"]
 
+[[test]]
+name = "ssl_conf_psk_cb"
+required-features = ["std"]
+
 [[test]]
 name = "ssl_conf_verify"
 required-features = ["std"]

mbedtls/src/ssl/config.rs

@@ -99,6 +99,7 @@ define!(
 callback!(DbgCallback: Fn(i32, Cow<'_, str>, i32, Cow<'_, str>) -> ());
 callback!(SniCallback: Fn(&mut HandshakeContext, &[u8]) -> Result<()>);
 callback!(CaCallback: Fn(&MbedtlsList<Certificate>) -> Result<MbedtlsList<Certificate>>);
+callback!(PskCallback: Fn(&mut HandshakeContext, &str) -> Result<()>);
 
 
 #[repr(transparent)]
@@ -164,6 +165,7 @@ define!(
         sni_callback: Option<Arc<dyn SniCallback + 'static>>,
         ticket_callback: Option<Arc<dyn TicketCallback + 'static>>,
         ca_callback: Option<Arc<dyn CaCallback + 'static>>,
+        psk_callback: Option<Arc<dyn PskCallback + 'static>>,
     };
     const drop: fn(&mut Self) = ssl_config_free;
     impl<'a> Into<ptr> {}
@@ -199,6 +201,7 @@ impl Config {
             sni_callback: None,
             ticket_callback: None,
             ca_callback: None,
+            psk_callback: None,
         }
     }
 
@@ -457,6 +460,43 @@ impl Config {
         self.dbg_callback = Some(Arc::new(cb));
         unsafe { ssl_conf_dbg(self.into(), Some(dbg_callback::<F>), &**self.dbg_callback.as_mut().unwrap() as *const _ as *mut c_void) }
     }
+
+    pub fn set_psk(&mut self, psk: &[u8], psk_identity: &str) -> Result<()> {
+        unsafe { ssl_conf_psk(&mut self.inner,
+                         psk.as_ptr(), psk.len(),
+                         psk_identity.as_ptr(), psk_identity.len())
+            .into_result().map(|_| ())
+        }
+    }
+
+    pub fn set_psk_callback<F>(&mut self, cb: F)
+        where
+            F: PskCallback + 'static,
+    {
+        unsafe extern "C" fn psk_callback<F>(
+            closure: *mut c_void,
+            ctx: *mut ssl_context,
+            psk_identity: *const c_uchar,
+            identity_len: size_t,
+        ) -> c_int
+        where
+              F: PskCallback + 'static,
+        {
+            let cb = &mut *(closure as *mut F);
+            let context = UnsafeFrom::from(ctx).unwrap();
+
+            let mut ctx = HandshakeContext::init(context);
+
+            let psk_identity = std::str::from_utf8_unchecked(from_raw_parts(psk_identity, identity_len));
+            match cb(&mut ctx, psk_identity) {
+                Ok(()) => 0,
+                Err(e) => e.to_int(),
+            }
+        }
+
+        self.psk_callback = Some(Arc::new(cb));
+        unsafe { ssl_conf_psk_cb(self.into(), Some(psk_callback::<F>), &**self.psk_callback.as_mut().unwrap() as *const _ as *mut c_void) }
+    }
 }
 
 // TODO
@@ -466,8 +506,6 @@ impl Config {
 // ssl_conf_dtls_badmac_limit
 // ssl_conf_handshake_timeout
 // ssl_conf_session_cache
-// ssl_conf_psk
-// ssl_conf_psk_cb
 // ssl_conf_sig_hashes
 // ssl_conf_alpn_protocols
 // ssl_conf_fallback

mbedtls/src/ssl/context.rs

@@ -408,6 +408,14 @@ impl HandshakeContext {
 
         Ok(())
     }
+
+    pub fn set_psk(&mut self, psk: &[u8]) -> Result<()> {
+        unsafe {
+            ssl_set_hs_psk(self.context.into(), psk.as_ptr(), psk.len()).into_result()?;
+        }
+
+        Ok(())
+    }
 }
 
 #[cfg(test)]
@@ -485,7 +493,6 @@ mod tests {
 // ssl_renegotiate
 // ssl_send_alert_message
 // ssl_set_client_transport_id
-// ssl_set_hs_psk
 // ssl_set_timer_cb
 //
 // ssl_handshake_step

mbedtls/tests/ssl_conf_psk_cb.rs

@@ -0,0 +1,61 @@
+#![allow(dead_code)]
+extern crate mbedtls;
+
+use std::net::TcpStream;
+use std::sync::Arc;
+
+use mbedtls::rng::CtrDrbg;
+use mbedtls::ssl::config::{Endpoint, Preset, Transport};
+use mbedtls::ssl::{Config, Context};
+use mbedtls::ssl::context::HandshakeContext;
+use mbedtls::Result as TlsResult;
+use mbedtls::ssl::config::PskCallback;
+
+mod support;
+use support::entropy::entropy_new;
+
+fn client(conn: TcpStream, psk: &[u8]) -> TlsResult<()>
+{
+    let entropy = Arc::new(entropy_new());
+    let rng = Arc::new(CtrDrbg::new(entropy, None)?);
+    let mut config = Config::new(Endpoint::Client, Transport::Stream, Preset::Default);
+    config.set_rng(rng);
+    config.set_psk(psk, "Client_identity")?;
+    let mut ctx = Context::new(Arc::new(config));
+    ctx.establish(conn, None).map(|_| ())
+}
+
+fn server<F>(conn: TcpStream, psk_callback: F) -> TlsResult<()>
+    where
+        F: PskCallback + Send + 'static,
+{
+    let entropy = Arc::new(entropy_new());
+    let rng = Arc::new(CtrDrbg::new(entropy, None)?);
+    let mut config = Config::new(Endpoint::Server, Transport::Stream, Preset::Default);
+    config.set_rng(rng);
+    config.set_psk_callback(psk_callback);
+    let mut ctx = Context::new(Arc::new(config));
+    ctx.establish(conn, None).map(|_| ())
+}
+
+#[cfg(unix)]
+mod test {
+    use super::*;
+    use std::thread;
+    use crate::support::net::create_tcp_pair;
+    use crate::support::keys;
+
+    #[test]
+    fn callback_standard_psk() {
+        let (c, s) = create_tcp_pair().unwrap();
+
+        let psk_callback =
+            |ctx: &mut HandshakeContext, _: &str| {
+                ctx.set_psk(keys::PRESHARED_KEY)
+            };
+        let c = thread::spawn(move || super::client(c, keys::PRESHARED_KEY).unwrap());
+        let s = thread::spawn(move || super::server(s, psk_callback).unwrap());
+        c.join().unwrap();
+        s.join().unwrap();
+    }
+}

mbedtls/tests/support/keys.rs

@@ -92,3 +92,7 @@ pub const ROOT_CA_KEY: &'static str = concat!(include_str!("./keys/ca.key"),"\0"
 pub const EXPIRED_CERT_SUBJECT: &'static str = "CN=ExpiredNode";
 pub const EXPIRED_CERT: &'static str = concat!(include_str!("./keys/expired.crt"),"\0");
 pub const EXPIRED_KEY: &'static str = concat!(include_str!("./keys/expired.key"),"\0");
+
+pub const PRESHARED_KEY: &'static [u8] = &[
+    234, 206, 151, 23, 219,  21,  71, 144,
+    107,  42,  23, 67, 249, 173, 182, 224 ];