Do some tweaks and cleanups (#60)

Order imports as: std, third-party, crate.

Use `try_ref_from_ptr!` to turn `out_n`'s into refs. This moves
unsafe blocks from later in the function to earlier in the function,
closer to where their safety properties are checked. Also, remove
pre-emptive zeroing of `out_n`'s. According to the conventions listed
in our README, there are no partial successes or partial failures,
so we shouldn't write to `out_n` if we early-return.

Remove `strong_count < 1` checks on Arc. Follow-up from #32.

Add some missing comments and fix some formatting.

Author: jsha

src/cipher.rs

@@ -1,18 +1,17 @@
-use std::{cmp::min, slice};
-
-use crate::error::rustls_result;
-use crate::{
-    ffi_panic_boundary, ffi_panic_boundary_generic, ffi_panic_boundary_unit, try_ref_from_ptr,
-};
 use libc::{c_char, size_t};
 use std::io::Cursor;
 use std::os::raw::c_ushort;
 use std::sync::Arc;
+use std::{cmp::min, slice};
 
+use rustls::sign::CertifiedKey;
 use rustls::{Certificate, PrivateKey};
 use rustls_pemfile::{certs, pkcs8_private_keys, rsa_private_keys};
 
-use rustls::sign::CertifiedKey;
+use crate::error::rustls_result;
+use crate::{
+    ffi_panic_boundary, ffi_panic_boundary_generic, ffi_panic_boundary_unit, try_ref_from_ptr,
+};
 use rustls_result::NullParameter;
 
 /// All SignatureScheme currently defined in rustls.
@@ -49,12 +48,11 @@ fn signature_scheme_name(n: u16) -> String {
 
 /// Get the name of a SignatureScheme, represented by the `scheme` short value,
 /// if known by the rustls library. For unknown schemes, this returns a string
-/// with the scheme value in hex notation.
-/// Mainly useful for debugging output.
+/// with the scheme value in hex notation. Mainly useful for debugging output.
+///
 /// The caller provides `buf` for holding the string and gives its size as `len`
 /// bytes. On return `out_n` carries the number of bytes copied into `buf`. The
 /// `buf` is not NUL-terminated.
-///
 #[no_mangle]
 pub extern "C" fn rustls_cipher_get_signature_scheme_name(
     scheme: c_ushort,
@@ -63,36 +61,46 @@ pub extern "C" fn rustls_cipher_get_signature_scheme_name(
     out_n: *mut size_t,
 ) {
     ffi_panic_boundary_unit! {
+        if buf.is_null() {
+            return;
+        }
         let write_buf: &mut [u8] = unsafe {
-            let out_n: &mut size_t = match out_n.as_mut() {
-                Some(out_n) => out_n,
-                None => return,
-            };
-            *out_n = 0;
-            if buf.is_null() {
-                return;
-            }
             slice::from_raw_parts_mut(buf as *mut u8, len as usize)
         };
+        let out_n: &mut size_t = try_ref_from_ptr!(out_n, &mut size_t, ());
         let name = signature_scheme_name(scheme);
         let len: usize = min(write_buf.len() - 1, name.len());
         write_buf[..len].copy_from_slice(&name.as_bytes()[..len]);
-        unsafe {
-            *out_n = len;
-        }
+        *out_n = len;
     }
 }
 
-/// The complete chain of certificates plus private key for
-/// being certified against someones list of trust anchors (commonly
-/// called root store). Corresponds to `CertifiedKey` in the Rust API.
+/// The complete chain of certificates to send during a TLS handshake,
+/// plus a private key that matches the end-entity (leaf) certificate.
+/// Corresponds to `CertifiedKey` in the Rust API.
+/// https://docs.rs/rustls/0.19.0/rustls/sign/struct.CertifiedKey.html
 pub struct rustls_certified_key {
     // We use the opaque struct pattern to tell C about our types without
     // telling them what's inside.
     // https://doc.rust-lang.org/nomicon/ffi.html#representing-opaque-structs
     _private: [u8; 0],
 }
 
+/// Build a `rustls_certified_key` from a certificate chain and a private key.
+/// `cert_chain` must point to a buffer of `cert_chain_len` bytes, containing
+/// a series of PEM-encoded certificates, with the end-entity (leaf)
+/// certificate first.
+///
+/// `private_key` must point to a buffer of `private_key_len` bytes, containing
+/// a PEM-encoded private key in either PKCS#1 or PKCS#8 format.
+///
+/// On success, this writes a pointer to the newly created
+/// `rustls_certified_key` in `certified_key_out`. That pointer must later
+/// be freed with `rustls_certified_key_free` to avoid memory leaks. Note that
+/// internally, this is an atomically reference-counted pointer, so even after
+/// the original caller has called `rustls_certified_key_free`, other objects
+/// may retain a pointer to the object. The memory will be freed when all
+/// references are gone.
 #[no_mangle]
 pub extern "C" fn rustls_certified_key_build(
     cert_chain: *const u8,
@@ -102,14 +110,15 @@ pub extern "C" fn rustls_certified_key_build(
     certified_key_out: *mut *const rustls_certified_key,
 ) -> rustls_result {
     ffi_panic_boundary! {
+        let certified_key_out: &mut *const rustls_certified_key =
+            try_ref_from_ptr!(certified_key_out, &mut *const rustls_certified_key);
         let certified_key = match certified_key_build(
             cert_chain, cert_chain_len, private_key, private_key_len) {
             Ok(key) => Box::new(key),
             Err(rr) => return rr,
         };
-        unsafe {
-            *certified_key_out = Arc::into_raw(Arc::new(*certified_key)) as *const _;
-        }
+        let certified_key = Arc::into_raw(Arc::new(*certified_key)) as *const _;
+        *certified_key_out = certified_key;
         return rustls_result::Ok
     }
 }
@@ -127,19 +136,11 @@ pub extern "C" fn rustls_certified_key_free(config: *const rustls_certified_key)
         // To free the certified_key, we reconstruct the Arc. It should have a refcount of 1,
         // representing the C code's copy. When it drops, that refcount will go down to 0
         // and the inner ServerConfig will be dropped.
-        let arc: Arc<CertifiedKey> = unsafe { Arc::from_raw(key) };
-        let strong_count = Arc::strong_count(&arc);
-        if strong_count < 1 {
-            eprintln!(
-                "rustls_certified_key_free: invariant failed: arc.strong_count was < 1: {}. \
-                You must not free the same certified_key multiple times.",
-                strong_count
-            );
-        }
+        unsafe { drop(Arc::from_raw(key)) };
     }
 }
 
-pub(crate) fn certified_key_build(
+fn certified_key_build(
     cert_chain: *const u8,
     cert_chain_len: size_t,
     private_key: *const u8,

src/client.rs

@@ -483,12 +483,7 @@ pub extern "C" fn rustls_client_session_read(
             }
             slice::from_raw_parts_mut(buf, count as usize)
         };
-        let out_n = unsafe {
-            match out_n.as_mut() {
-                Some(out_n) => out_n,
-                None => return NullParameter,
-            }
-        };
+        let out_n: &mut size_t = try_ref_from_ptr!(out_n, &mut size_t);
         let n_read: usize = match session.read(read_buf) {
             Ok(n) => n,
             // The CloseNotify TLS alert is benign, but rustls returns it as an Error. See comment on
@@ -528,12 +523,7 @@ pub extern "C" fn rustls_client_session_read_tls(
             }
             slice::from_raw_parts(buf, count as usize)
         };
-        let out_n = unsafe {
-            match out_n.as_mut() {
-                Some(out_n) => out_n,
-                None => return NullParameter,
-            }
-        };
+        let out_n: &mut size_t = try_ref_from_ptr!(out_n, &mut size_t);
         let mut cursor = Cursor::new(input_buf);
         let n_read: usize = match session.read_tls(&mut cursor) {
             Ok(n) => n,
@@ -569,12 +559,7 @@ pub extern "C" fn rustls_client_session_write_tls(
             }
             slice::from_raw_parts_mut(buf, count as usize)
         };
-        let out_n = unsafe {
-            match out_n.as_mut() {
-                Some(out_n) => out_n,
-                None => return NullParameter,
-            }
-        };
+        let out_n: &mut size_t = try_ref_from_ptr!(out_n, &mut size_t);
         let n_written: usize = match session.write_tls(&mut output_buf) {
             Ok(n) => n,
             Err(_) => return rustls_result::Io,