feat(SHA): Borrow the SHA peripheral instead of unsafely stealing

Author: AnthonyGrondin

esp-mbedtls/src/compat/edge_nal_compat.rs

@@ -4,10 +4,13 @@ use core::{
     cell::{Cell, RefCell, UnsafeCell},
     mem::MaybeUninit,
     net::{Ipv4Addr, SocketAddr, SocketAddrV4},
+    ops::DerefMut,
     ptr::NonNull,
 };
 
+use crate::{hal::peripheral::PeripheralRef, SHA};
 use edge_nal::{Close, TcpBind};
+
 use edge_nal_embassy::{Tcp, TcpAccept, TcpSocket};
 
 pub struct TlsAcceptor<
@@ -23,6 +26,7 @@ pub struct TlsAcceptor<
     owns_rsa: bool,
     tls_buffers: &'d TlsBuffers<RX_SZ, TX_SZ>,
     tls_buffers_ptr: RefCell<NonNull<([u8; RX_SZ], [u8; TX_SZ])>>,
+    sha: RefCell<PeripheralRef<'d, SHA>>,
 }
 
 impl<'d, D, const N: usize, const RX_SZ: usize, const TX_SZ: usize> Drop
@@ -55,6 +59,7 @@ where
         port: u16,
         version: TlsVersion,
         certificates: Certificates<'d>,
+        sha: impl Peripheral<P = SHA> + 'd,
     ) -> Self {
         let acceptor = tcp
             .bind(SocketAddr::V4(SocketAddrV4::new(
@@ -73,6 +78,7 @@ where
             owns_rsa: false,
             tls_buffers,
             tls_buffers_ptr: RefCell::new(socket_buffers),
+            sha: sha.into_ref().into(),
         }
     }
 
@@ -146,6 +152,7 @@ where
             self.certificates,
             rx,
             tx,
+            self.sha.borrow_mut().reborrow().deref_mut(),
         )?;
 
         log::debug!("Establishing SSL connection");

esp-mbedtls/src/lib.rs

@@ -495,10 +495,13 @@ impl<T> Session<T> {
         mode: Mode,
         min_version: TlsVersion,
         certificates: Certificates,
+        sha: impl Peripheral<P = SHA>,
     ) -> Result<Self, TlsError> {
-        // TODO: Take peripheral from user
-        let sha = Sha::new(unsafe { SHA::steal() });
-        critical_section::with(|cs| SHARED_SHA.borrow_ref_mut(cs).replace(sha));
+        critical_section::with(|cs| {
+            SHARED_SHA
+                .borrow_ref_mut(cs)
+                .replace(unsafe { core::mem::transmute(Sha::new(sha)) })
+        });
 
         let (drbg_context, ssl_context, ssl_config, crt, client_crt, private_key) =
             certificates.init_ssl(servername, mode, min_version)?;
@@ -641,6 +644,7 @@ impl<T> Drop for Session<T> {
             if self.owns_rsa {
                 RSA_REF = core::mem::transmute(None::<RSA>);
             }
+            critical_section::with(|cs| SHARED_SHA.borrow_ref_mut(cs).take());
             mbedtls_ssl_close_notify(self.ssl_context);
             mbedtls_ctr_drbg_free(self.drbg_context);
             mbedtls_ssl_config_free(self.ssl_config);
@@ -751,10 +755,13 @@ pub mod asynch {
 
             rx_buffer: &'a mut [u8; RX_SIZE],
             tx_buffer: &'a mut [u8; TX_SIZE],
+            sha: impl Peripheral<P = SHA>,
         ) -> Result<Self, TlsError> {
-            // TODO: Take peripheral from user
-            let sha = Sha::new(unsafe { SHA::steal() });
-            critical_section::with(|cs| SHARED_SHA.borrow_ref_mut(cs).replace(sha));
+            critical_section::with(|cs| {
+                SHARED_SHA
+                    .borrow_ref_mut(cs)
+                    .replace(unsafe { core::mem::transmute(Sha::new(sha)) })
+            });
 
             let (drbg_context, ssl_context, ssl_config, crt, client_crt, private_key) =
                 certificates.init_ssl(servername, mode, min_version)?;
@@ -797,6 +804,7 @@ pub mod asynch {
                 if self.owns_rsa {
                     RSA_REF = core::mem::transmute(None::<RSA>);
                 }
+                critical_section::with(|cs| SHARED_SHA.borrow_ref_mut(cs).take());
                 mbedtls_ssl_close_notify(self.ssl_context);
                 mbedtls_ctr_drbg_free(self.drbg_context);
                 mbedtls_ssl_config_free(self.ssl_config);

esp-mbedtls/src/sha/sha1.rs

@@ -10,8 +10,6 @@ pub struct mbedtls_sha1_context {
 
 #[no_mangle]
 pub unsafe extern "C" fn mbedtls_sha1_init(ctx: *mut mbedtls_sha1_context) {
-    let sha = crate::Sha::new(unsafe { crate::SHA::steal() });
-    critical_section::with(|cs| SHARED_SHA.borrow_ref_mut(cs).replace(sha));
     let hasher_mem =
         crate::calloc(1, core::mem::size_of::<Context<Sha1>>() as u32) as *mut Context<Sha1>;
     core::ptr::write(hasher_mem, Context::<Sha1>::new());

examples/async_client.rs

@@ -140,6 +140,7 @@ async fn main(spawner: Spawner) -> ! {
         },
         mk_static!([u8; 4096], [0; 4096]),
         mk_static!([u8; 4096], [0; 4096]),
+        peripherals.SHA,
     )
     .unwrap()
     .with_hardware_rsa(peripherals.RSA);

examples/async_client_mTLS.rs

@@ -146,6 +146,7 @@ async fn main(spawner: Spawner) -> ! {
         certificates,
         mk_static!([u8; 4096], [0; 4096]),
         mk_static!([u8; 4096], [0; 4096]),
+        peripherals.SHA,
     )
     .unwrap()
     .with_hardware_rsa(peripherals.RSA);

examples/async_server.rs

@@ -164,6 +164,7 @@ async fn main(spawner: Spawner) -> ! {
             },
             tls_rx_buffer,
             tls_tx_buffer,
+            &mut peripherals.SHA,
         )
         .unwrap()
         .with_hardware_rsa(&mut peripherals.RSA);

examples/async_server_mTLS.rs

@@ -182,6 +182,7 @@ async fn main(spawner: Spawner) -> ! {
             },
             tls_rx_buffer,
             tls_tx_buffer,
+            &mut peripherals.SHA,
         )
         .unwrap()
         .with_hardware_rsa(&mut peripherals.RSA);

examples/edge_server.rs

@@ -34,7 +34,7 @@ use esp_wifi::wifi::{
     WifiState,
 };
 use esp_wifi::{init, EspWifiInitFor};
-use hal::{prelude::*, rng::Rng, timer::timg::TimerGroup};
+use hal::{peripherals::SHA, prelude::*, rng::Rng, timer::timg::TimerGroup};
 
 // Patch until https://github.com/embassy-rs/static-cell/issues/16 is fixed
 macro_rules! mk_static {
@@ -164,13 +164,16 @@ async fn main(spawner: Spawner) -> ! {
         ..Default::default()
     };
 
+    let sha = mk_static!(SHA, peripherals.SHA);
+
     loop {
         let tls_acceptor = esp_mbedtls::asynch::TlsAcceptor::new(
             tcp,
             tls_buffers,
             443,
             TlsVersion::Tls1_2,
             certificates,
+            &mut *sha,
         )
         .await
         .with_hardware_rsa(&mut peripherals.RSA);

examples/sync_client.rs

@@ -116,6 +116,7 @@ fn main() -> ! {
             .ok(),
             ..Default::default()
         },
+        peripherals.SHA,
     )
     .unwrap()
     .with_hardware_rsa(peripherals.RSA);

examples/sync_client_mTLS.rs

@@ -123,6 +123,7 @@ fn main() -> ! {
         Mode::Client,
         TlsVersion::Tls1_3,
         certificates,
+        peripherals.SHA,
     )
     .unwrap()
     .with_hardware_rsa(peripherals.RSA);

examples/sync_server.rs

@@ -137,6 +137,7 @@ fn main() -> ! {
                     .ok(),
                     ..Default::default()
                 },
+                &mut peripherals.SHA,
             )
             .unwrap()
             .with_hardware_rsa(&mut peripherals.RSA);

examples/sync_server_mTLS.rs

@@ -158,6 +158,7 @@ fn main() -> ! {
                     .ok(),
                     ..Default::default()
                 },
+                &mut peripherals.SHA,
             )
             .unwrap()
             .with_hardware_rsa(&mut peripherals.RSA);