// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{
    callbacks::*,
    enums::*,
    error::{Error, Fallible},
    security,
};
use core::{convert::TryInto, ptr::NonNull};
use s2n_tls_sys::*;
use std::{
    ffi::{c_void, CString},
    path::Path,
    sync::atomic::{AtomicUsize, Ordering},
};

#[derive(Debug, PartialEq)]
pub struct Config(NonNull<s2n_config>);

/// # Safety
///
/// Safety: s2n_config objects can be sent across threads
unsafe impl Send for Config {}

/// # Safety
///
/// Safety: All C methods that mutate the s2n_config are wrapped
/// in Rust methods that require a mutable reference.
unsafe impl Sync for Config {}

impl Config {
    /// Returns a Config object with pre-defined defaults.
    ///
    /// Use the [`Builder`] if custom configuration is desired.
    pub fn new() -> Self {
        Self::default()
    }

    /// Returns a Builder which can be used to configure the Config
    pub fn builder() -> Builder {
        Builder::default()
    }

    /// # Safety
    ///
    /// This config _MUST_ have been initialized with a [`Builder`].
    /// Additionally, this does NOT increment the config reference count,
    /// so consider cloning the result if the source pointer is still
    /// valid and usable afterwards.
    pub(crate) unsafe fn from_raw(config: NonNull<s2n_config>) -> Self {
        let config = Self(config);

        // Check if the context can be retrieved.
        // If it can't, this is not a valid config.
        config.context();

        config
    }

    pub(crate) fn as_mut_ptr(&mut self) -> *mut s2n_config {
        self.0.as_ptr()
    }

    /// Retrieve a reference to the [`Context`] stored on the config.
    pub(crate) fn context(&self) -> &Context {
        let mut ctx = core::ptr::null_mut();
        unsafe {
            s2n_config_get_ctx(self.0.as_ptr(), &mut ctx)
                .into_result()
                .unwrap();
            &*(ctx as *const Context)
        }
    }

    /// Retrieve a mutable reference to the [`Context`] stored on the config.
    pub(crate) fn context_mut(&mut self) -> &mut Context {
        let mut ctx = core::ptr::null_mut();
        unsafe {
            s2n_config_get_ctx(self.as_mut_ptr(), &mut ctx)
                .into_result()
                .unwrap();
            &mut *(ctx as *mut Context)
        }
    }

    #[cfg(test)]
    /// Get the refcount associated with the config
    pub fn test_get_refcount(&self) -> Result<usize, Error> {
        let context = self.context();
        Ok(context.refcount.load(Ordering::SeqCst))
    }
}

impl Default for Config {
    fn default() -> Self {
        Builder::new().build().unwrap()
    }
}

impl Clone for Config {
    fn clone(&self) -> Self {
        let context = self.context();

        // Safety
        //
        // Using a relaxed ordering is alright here, as knowledge of the
        // original reference prevents other threads from erroneously deleting
        // the object.
        // https://github.com/rust-lang/rust/blob/e012a191d768adeda1ee36a99ef8b92d51920154/library/alloc/src/sync.rs#L1329
        let _count = context.refcount.fetch_add(1, Ordering::Relaxed);
        Self(self.0)
    }
}

impl Drop for Config {
    fn drop(&mut self) {
        let context = self.context_mut();
        let count = context.refcount.fetch_sub(1, Ordering::Release);
        debug_assert!(count > 0, "refcount should not drop below 1 instance");

        // only free the config if this is the last instance
        if count != 1 {
            return;
        }

        // Safety
        //
        // The use of Ordering and fence mirrors the `Arc` implementation in
        // the standard library.
        //
        // This fence is needed to prevent reordering of use of the data and
        // deletion of the data.  Because it is marked `Release`, the decreasing
        // of the reference count synchronizes with this `Acquire` fence. This
        // means that use of the data happens before decreasing the reference
        // count, which happens before this fence, which happens before the
        // deletion of the data.
        // https://github.com/rust-lang/rust/blob/e012a191d768adeda1ee36a99ef8b92d51920154/library/alloc/src/sync.rs#L1637
        std::sync::atomic::fence(Ordering::Acquire);

        unsafe {
            // This is the last instance so free the context.
            let context = Box::from_raw(context);
            drop(context);

            let _ = s2n_config_free(self.0.as_ptr()).into_result();
        }
    }
}

#[derive(Default)]
pub struct Builder {
    config: Config,
    load_system_certs: bool,
    enable_ocsp: bool,
}

impl Builder {
    pub fn new() -> Self {
        crate::init::init();
        let config = unsafe { s2n_config_new_minimal().into_result() }.unwrap();

        let context = Box::<Context>::default();
        let context = Box::into_raw(context) as *mut c_void;

        unsafe {
            s2n_config_set_ctx(config.as_ptr(), context)
                .into_result()
                .unwrap();

            // The client hello callback originally did not support async operations,
            // so defaults to blocking mode for backwards compatibility with old integrations.
            // But these bindings use a polling model, so assume non-blocking mode.
            s2n_config_set_client_hello_cb_mode(
                config.as_ptr(),
                s2n_client_hello_cb_mode::NONBLOCKING,
            )
            .into_result()
            .unwrap();
        }

        Self {
            config: Config(config),
            load_system_certs: true,
            enable_ocsp: false,
        }
    }

    pub fn set_alert_behavior(&mut self, value: AlertBehavior) -> Result<&mut Self, Error> {
        unsafe { s2n_config_set_alert_behavior(self.as_mut_ptr(), value.into()).into_result() }?;
        Ok(self)
    }

    pub fn set_security_policy(&mut self, policy: &security::Policy) -> Result<&mut Self, Error> {
        unsafe {
            s2n_config_set_cipher_preferences(self.as_mut_ptr(), policy.as_cstr().as_ptr())
                .into_result()
        }?;
        Ok(self)
    }

    /// sets the application protocol preferences on an s2n_config object.
    ///
    /// protocols is a list in order of preference, with most preferred protocol first,
    /// and of length protocol_count. When acting as a client the protocol list is
    /// included in the Client Hello message as the ALPN extension. As a server, the
    /// list is used to negotiate a mutual application protocol with the client. After
    /// the negotiation for the connection has completed, the agreed upon protocol can
    /// be retrieved with s2n_get_application_protocol
    pub fn set_application_protocol_preference<P: IntoIterator<Item = I>, I: AsRef<[u8]>>(
        &mut self,
        protocols: P,
    ) -> Result<&mut Self, Error> {
        // reset the list
        unsafe {
            s2n_config_set_protocol_preferences(self.as_mut_ptr(), core::ptr::null(), 0)
                .into_result()
        }?;

        for protocol in protocols {
            self.append_application_protocol_preference(protocol.as_ref())?;
        }

        Ok(self)
    }

    pub fn append_application_protocol_preference(
        &mut self,
        protocol: &[u8],
    ) -> Result<&mut Self, Error> {
        unsafe {
            s2n_config_append_protocol_preference(
                self.as_mut_ptr(),
                protocol.as_ptr(),
                protocol
                    .len()
                    .try_into()
                    .map_err(|_| Error::INVALID_INPUT)?,
            )
            .into_result()
        }?;
        Ok(self)
    }

    /// Turns off x509 verification
    ///
    /// # Safety
    /// This functionality will weaken the security of the connections. As such, it should only
    /// be used in development environments where obtaining a valid certificate would not be possible.
    pub unsafe fn disable_x509_verification(&mut self) -> Result<&mut Self, Error> {
        s2n_config_disable_x509_verification(self.as_mut_ptr()).into_result()?;
        Ok(self)
    }

    pub fn add_dhparams(&mut self, pem: &[u8]) -> Result<&mut Self, Error> {
        let cstring = CString::new(pem).map_err(|_| Error::INVALID_INPUT)?;
        unsafe { s2n_config_add_dhparams(self.as_mut_ptr(), cstring.as_ptr()).into_result() }?;
        Ok(self)
    }

    pub fn load_pem(&mut self, certificate: &[u8], private_key: &[u8]) -> Result<&mut Self, Error> {
        let certificate = CString::new(certificate).map_err(|_| Error::INVALID_INPUT)?;
        let private_key = CString::new(private_key).map_err(|_| Error::INVALID_INPUT)?;
        unsafe {
            s2n_config_add_cert_chain_and_key(
                self.as_mut_ptr(),
                certificate.as_ptr(),
                private_key.as_ptr(),
            )
            .into_result()
        }?;
        Ok(self)
    }

    pub fn load_public_pem(&mut self, certificate: &[u8]) -> Result<&mut Self, Error> {
        let size: u32 = certificate
            .len()
            .try_into()
            .map_err(|_| Error::INVALID_INPUT)?;
        let certificate = certificate.as_ptr() as *mut u8;
        unsafe { s2n_config_add_cert_chain(self.as_mut_ptr(), certificate, size) }.into_result()?;
        Ok(self)
    }

    pub fn trust_pem(&mut self, certificate: &[u8]) -> Result<&mut Self, Error> {
        let certificate = CString::new(certificate).map_err(|_| Error::INVALID_INPUT)?;
        unsafe {
            s2n_config_add_pem_to_trust_store(self.as_mut_ptr(), certificate.as_ptr()).into_result()
        }?;
        Ok(self)
    }

    /// Adds to the trust store from a CA file or directory containing trusted certificates.
    ///
    /// NOTE: This function is equivalent to `s2n_config_set_verification_ca_location` except it does
    /// not automatically enable the client to request OCSP stapling from the server.
    pub fn trust_location(
        &mut self,
        file: Option<&Path>,
        dir: Option<&Path>,
    ) -> Result<&mut Self, Error> {
        fn to_cstr(input: Option<&Path>) -> Result<Option<CString>, Error> {
            Ok(match input {
                Some(input) => {
                    let string = input.to_str().ok_or(Error::INVALID_INPUT)?;
                    let cstring = CString::new(string).map_err(|_| Error::INVALID_INPUT)?;
                    Some(cstring)
                }
                None => None,
            })
        }

        let file_cstr = to_cstr(file)?;
        let file_ptr = file_cstr
            .as_ref()
            .map(|f| f.as_ptr())
            .unwrap_or(core::ptr::null());

        let dir_cstr = to_cstr(dir)?;
        let dir_ptr = dir_cstr
            .as_ref()
            .map(|f| f.as_ptr())
            .unwrap_or(core::ptr::null());

        unsafe {
            s2n_config_set_verification_ca_location(self.as_mut_ptr(), file_ptr, dir_ptr)
                .into_result()
        }?;

        // If OCSP has not been explicitly requested, turn off OCSP. This is to prevent this function from
        // automatically enabling `OCSP` due to the legacy behavior of `s2n_config_set_verification_ca_location`
        if !self.enable_ocsp {
            unsafe {
                s2n_config_set_status_request_type(self.as_mut_ptr(), s2n_status_request_type::NONE)
                    .into_result()?
            };
        }

        Ok(self)
    }

    /// Sets whether or not default system certificates will be loaded into the trust store.
    ///
    /// Set to false for increased performance if system certificates are not needed during
    /// certificate validation.
    pub fn with_system_certs(&mut self, load_system_certs: bool) -> Result<&mut Self, Error> {
        self.load_system_certs = load_system_certs;
        Ok(self)
    }

    pub fn wipe_trust_store(&mut self) -> Result<&mut Self, Error> {
        unsafe { s2n_config_wipe_trust_store(self.as_mut_ptr()).into_result()? };
        Ok(self)
    }

    /// Sets whether or not a client certificate should be required to complete the TLS connection.
    ///
    /// See the [Usage Guide](https://github.com/aws/s2n-tls/blob/main/docs/USAGE-GUIDE.md#client-auth-related-calls) for more details.
    pub fn set_client_auth_type(&mut self, auth_type: ClientAuthType) -> Result<&mut Self, Error> {
        unsafe {
            s2n_config_set_client_auth_type(self.as_mut_ptr(), auth_type.into()).into_result()
        }?;
        Ok(self)
    }

    /// Clients will request OCSP stapling from the server.
    pub fn enable_ocsp(&mut self) -> Result<&mut Self, Error> {
        unsafe {
            s2n_config_set_status_request_type(self.as_mut_ptr(), s2n_status_request_type::OCSP)
                .into_result()
        }?;
        self.enable_ocsp = true;
        Ok(self)
    }

    /// Sets the OCSP data for the default certificate chain associated with the Config.
    ///
    /// Servers will send the data in response to OCSP stapling requests from clients.
    //
    // NOTE: this modifies a certificate chain, NOT the Config itself. This is currently safe
    // because the certificate chain is set with s2n_config_add_cert_chain_and_key, which
    // creates a new certificate chain only accessible by the given config. It will
    // NOT be safe when we add support for the newer s2n_config_add_cert_chain_and_key_to_store API,
    // which allows certificate chains to be shared across configs.
    // In that case, we'll need additional guard rails either in these bindings or in the underlying C.
    pub fn set_ocsp_data(&mut self, data: &[u8]) -> Result<&mut Self, Error> {
        let size: u32 = data.len().try_into().map_err(|_| Error::INVALID_INPUT)?;
        unsafe {
            s2n_config_set_extension_data(
                self.as_mut_ptr(),
                s2n_tls_extension_type::OCSP_STAPLING,
                data.as_ptr(),
                size,
            )
            .into_result()
        }?;
        self.enable_ocsp()
    }

    /// Sets the callback to use for verifying that a hostname from an X.509 certificate is
    /// trusted.
    ///
    /// The callback may be called more than once during certificate validation as each SAN on
    /// the certificate will be checked.
    ///
    /// Corresponds to the underlying C API
    /// [s2n_config_set_verify_host_callback](https://aws.github.io/s2n-tls/doxygen/s2n_8h.html).
    pub fn set_verify_host_callback<T: 'static + VerifyHostNameCallback>(
        &mut self,
        handler: T,
    ) -> Result<&mut Self, Error> {
        unsafe extern "C" fn verify_host_cb_fn(
            host_name: *const ::libc::c_char,
            host_name_len: usize,
            context: *mut ::libc::c_void,
        ) -> u8 {
            let context = &mut *(context as *mut Context);
            let handler = context.verify_host_callback.as_mut().unwrap();
            verify_host(host_name, host_name_len, handler)
        }

        self.config.context_mut().verify_host_callback = Some(Box::new(handler));
        unsafe {
            s2n_config_set_verify_host_callback(
                self.as_mut_ptr(),
                Some(verify_host_cb_fn),
                self.config.context_mut() as *mut Context as *mut c_void,
            )
            .into_result()?;
        }
        Ok(self)
    }

    /// # Safety
    /// THIS SHOULD BE USED FOR DEBUGGING PURPOSES ONLY!
    /// The `context` pointer must live at least as long as the config
    pub unsafe fn set_key_log_callback(
        &mut self,
        callback: s2n_key_log_fn,
        context: *mut core::ffi::c_void,
    ) -> Result<&mut Self, Error> {
        s2n_config_set_key_log_cb(self.as_mut_ptr(), callback, context).into_result()?;
        Ok(self)
    }

    pub fn set_max_cert_chain_depth(&mut self, depth: u16) -> Result<&mut Self, Error> {
        unsafe { s2n_config_set_max_cert_chain_depth(self.as_mut_ptr(), depth).into_result() }?;
        Ok(self)
    }

    pub fn set_send_buffer_size(&mut self, size: u32) -> Result<&mut Self, Error> {
        unsafe { s2n_config_set_send_buffer_size(self.as_mut_ptr(), size).into_result() }?;
        Ok(self)
    }

    /// Set a custom callback function which is run after parsing the client hello.
    pub fn set_client_hello_callback<T: 'static + ClientHelloCallback>(
        &mut self,
        handler: T,
    ) -> Result<&mut Self, Error> {
        unsafe extern "C" fn client_hello_cb(
            connection_ptr: *mut s2n_connection,
            _context: *mut core::ffi::c_void,
        ) -> libc::c_int {
            with_context(connection_ptr, |conn, context| {
                let callback = context.client_hello_callback.as_ref();
                let future = callback
                    .map(|c| c.on_client_hello(conn))
                    .unwrap_or(Ok(None));
                AsyncCallback::trigger_client_hello_cb(future, conn)
            })
            .into()
        }

        let handler = Box::new(handler);
        let context = self.config.context_mut();
        context.client_hello_callback = Some(handler);

        unsafe {
            s2n_config_set_client_hello_cb(
                self.as_mut_ptr(),
                Some(client_hello_cb),
                core::ptr::null_mut(),
            )
            .into_result()?;
        }

        Ok(self)
    }

    /// Set a callback function triggered by operations requiring the private key.
    ///
    /// See https://github.com/aws/s2n-tls/blob/main/docs/USAGE-GUIDE.md#private-key-operation-related-calls
    pub fn set_private_key_callback<T: 'static + PrivateKeyCallback>(
        &mut self,
        handler: T,
    ) -> Result<&mut Self, Error> {
        unsafe extern "C" fn private_key_cb(
            conn_ptr: *mut s2n_connection,
            op_ptr: *mut s2n_async_pkey_op,
        ) -> libc::c_int {
            with_context(conn_ptr, |conn, context| {
                let state = PrivateKeyOperation::try_from_cb(conn, op_ptr);
                let callback = context.private_key_callback.as_ref();
                let future_result = state.and_then(|state| {
                    callback.map_or(Ok(None), |callback| callback.handle_operation(conn, state))
                });
                AsyncCallback::trigger(future_result, conn)
            })
            .into()
        }

        let handler = Box::new(handler);
        let context = self.config.context_mut();
        context.private_key_callback = Some(handler);

        unsafe {
            s2n_config_set_async_pkey_callback(self.as_mut_ptr(), Some(private_key_cb))
                .into_result()?;
        }
        Ok(self)
    }

    /// Set a callback function that will be used to get the system time.
    ///
    /// The wall clock time is the best-guess at the real time, measured since the epoch.
    /// Unlike monotonic time, it CAN move backwards.
    /// It is used by s2n-tls for timestamps.
    pub fn set_wall_clock<T: 'static + WallClock>(
        &mut self,
        handler: T,
    ) -> Result<&mut Self, Error> {
        unsafe extern "C" fn clock_cb(
            context: *mut ::libc::c_void,
            time_in_nanos: *mut u64,
        ) -> libc::c_int {
            let context = &mut *(context as *mut Context);
            if let Some(handler) = context.wall_clock.as_mut() {
                if let Ok(nanos) = handler.get_time_since_epoch().as_nanos().try_into() {
                    *time_in_nanos = nanos;
                    return CallbackResult::Success.into();
                }
            }
            CallbackResult::Failure.into()
        }

        let handler = Box::new(handler);
        let context = self.config.context_mut();
        context.wall_clock = Some(handler);
        unsafe {
            s2n_config_set_wall_clock(
                self.as_mut_ptr(),
                Some(clock_cb),
                self.config.context_mut() as *mut _ as *mut c_void,
            )
            .into_result()?;
        }
        Ok(self)
    }

    /// Set a callback function that will be used to get the monotonic time.
    ///
    /// The monotonic time is the time since an arbitrary, unspecified point.
    /// Unlike wall clock time, it MUST never move backwards.
    /// It is used by s2n-tls for timers.
    pub fn set_monotonic_clock<T: 'static + MonotonicClock>(
        &mut self,
        handler: T,
    ) -> Result<&mut Self, Error> {
        unsafe extern "C" fn clock_cb(
            context: *mut ::libc::c_void,
            time_in_nanos: *mut u64,
        ) -> libc::c_int {
            let context = &mut *(context as *mut Context);
            if let Some(handler) = context.monotonic_clock.as_mut() {
                if let Ok(nanos) = handler.get_time().as_nanos().try_into() {
                    *time_in_nanos = nanos;
                    return CallbackResult::Success.into();
                }
            }
            CallbackResult::Failure.into()
        }

        let handler = Box::new(handler);
        let context = self.config.context_mut();
        context.monotonic_clock = Some(handler);
        unsafe {
            s2n_config_set_monotonic_clock(
                self.as_mut_ptr(),
                Some(clock_cb),
                self.config.context_mut() as *mut _ as *mut c_void,
            )
            .into_result()?;
        }
        Ok(self)
    }

    pub fn build(mut self) -> Result<Config, Error> {
        if self.load_system_certs {
            unsafe {
                s2n_config_load_system_certs(self.as_mut_ptr()).into_result()?;
            }
        }

        Ok(self.config)
    }

    fn as_mut_ptr(&mut self) -> *mut s2n_config {
        self.config.as_mut_ptr()
    }
}

#[cfg(feature = "quic")]
impl Builder {
    pub fn enable_quic(&mut self) -> Result<&mut Self, Error> {
        unsafe { s2n_tls_sys::s2n_config_enable_quic(self.as_mut_ptr()).into_result() }?;
        Ok(self)
    }
}

pub(crate) struct Context {
    refcount: AtomicUsize,
    pub(crate) client_hello_callback: Option<Box<dyn ClientHelloCallback>>,
    pub(crate) private_key_callback: Option<Box<dyn PrivateKeyCallback>>,
    pub(crate) verify_host_callback: Option<Box<dyn VerifyHostNameCallback>>,
    pub(crate) wall_clock: Option<Box<dyn WallClock>>,
    pub(crate) monotonic_clock: Option<Box<dyn MonotonicClock>>,
}

impl Default for Context {
    fn default() -> Self {
        // The AtomicUsize is used to manually track the reference count of the Config.
        // This mechanism is used to track when the Config object should be freed.
        let refcount = AtomicUsize::new(1);

        Self {
            refcount,
            client_hello_callback: None,
            private_key_callback: None,
            verify_host_callback: None,
            wall_clock: None,
            monotonic_clock: None,
        }
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

#![allow(clippy::missing_safety_doc)] // TODO add safety docs

use crate::{
    callbacks::*,
    config::Config,
    enums::*,
    error::{Error, Fallible, Pollable},
    security,
};

use core::{
    convert::TryInto,
    fmt,
    mem::{self, ManuallyDrop, MaybeUninit},
    pin::Pin,
    ptr::NonNull,
    task::{Poll, Waker},
    time::Duration,
};
use libc::c_void;
use s2n_tls_sys::*;
use std::ffi::CStr;

mod builder;
pub use builder::*;

macro_rules! static_const_str {
    ($c_chars:expr) => {
        unsafe { CStr::from_ptr($c_chars) }
            .to_str()
            .map_err(|_| Error::INVALID_INPUT)
    };
}

pub struct Connection {
    connection: NonNull<s2n_connection>,
}

impl fmt::Debug for Connection {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut debug = f.debug_struct("Connection");
        if let Ok(handshake) = self.handshake_type() {
            debug.field("handshake_type", &handshake);
        }
        if let Ok(cipher) = self.cipher_suite() {
            debug.field("cipher_suite", &cipher);
        }
        if let Ok(version) = self.actual_protocol_version() {
            debug.field("actual_protocol_version", &version);
        }
        if let Ok(curve) = self.selected_curve() {
            debug.field("selected_curve", &curve);
        }
        debug.finish_non_exhaustive()
    }
}

/// # Safety
///
/// s2n_connection objects can be sent across threads
unsafe impl Send for Connection {}

/// # Safety
///
/// All C methods that mutate the s2n_connection are wrapped
/// in Rust methods that require a mutable reference.
unsafe impl Sync for Connection {}

impl Connection {
    pub fn new(mode: Mode) -> Self {
        crate::init::init();

        let connection = unsafe { s2n_connection_new(mode.into()).into_result() }.unwrap();

        unsafe {
            debug_assert! {
                s2n_connection_get_config(connection.as_ptr(), &mut core::ptr::null_mut())
                    .into_result()
                    .is_err()
            }
        }

        let mut connection = Self { connection };
        connection.init_context(mode);
        connection
    }

    fn init_context(&mut self, mode: Mode) {
        let context = Box::new(Context::new(mode));
        let context = Box::into_raw(context) as *mut c_void;
        // allocate a new context object
        unsafe {
            // There should never be an existing context
            debug_assert!(s2n_connection_get_ctx(self.connection.as_ptr())
                .into_result()
                .is_err());

            s2n_connection_set_ctx(self.connection.as_ptr(), context)
                .into_result()
                .unwrap();
        }
    }

    pub fn new_client() -> Self {
        Self::new(Mode::Client)
    }

    pub fn new_server() -> Self {
        Self::new(Mode::Server)
    }

    pub(crate) fn as_ptr(&mut self) -> *mut s2n_connection {
        self.connection.as_ptr()
    }

    /// # Safety
    ///
    /// Caller must ensure s2n_connection is a valid reference to a [`s2n_connection`] object
    pub(crate) unsafe fn from_raw(connection: NonNull<s2n_connection>) -> Self {
        Self { connection }
    }

    pub(crate) fn mode(&self) -> Mode {
        self.context().mode
    }

    /// can be used to configure s2n to either use built-in blinding (set blinding
    /// to Blinding::BuiltIn) or self-service blinding (set blinding to
    /// Blinding::SelfService).
    pub fn set_blinding(&mut self, blinding: Blinding) -> Result<&mut Self, Error> {
        unsafe {
            s2n_connection_set_blinding(self.connection.as_ptr(), blinding.into()).into_result()
        }?;
        Ok(self)
    }

    /// Reports the remaining nanoseconds before the connection may be gracefully shutdown.
    ///
    /// This method is expected to succeed, but could fail if the
    /// [underlying C call](`s2n_connection_get_delay`) encounters errors.
    /// Failure indicates that calls to [`Self::poll_shutdown`] will also fail and
    /// that a graceful two-way shutdown of the connection will not be possible.
    pub fn remaining_blinding_delay(&self) -> Result<Duration, Error> {
        let nanos = unsafe { s2n_connection_get_delay(self.connection.as_ptr()).into_result() }?;
        Ok(Duration::from_nanos(nanos))
    }

    /// Sets whether or not a Client Certificate should be required to complete the TLS Connection.
    ///
    /// If this is set to ClientAuthType::Optional the server will request a client certificate
    /// but allow the client to not provide one. Rejecting a client certificate when using
    /// ClientAuthType::Optional will terminate the handshake.
    pub fn set_client_auth_type(
        &mut self,
        client_auth_type: ClientAuthType,
    ) -> Result<&mut Self, Error> {
        unsafe {
            s2n_connection_set_client_auth_type(self.connection.as_ptr(), client_auth_type.into())
                .into_result()
        }?;
        Ok(self)
    }

    /// Attempts to drop the config on the connection.
    ///
    /// # Safety
    ///
    /// The caller must ensure the config associated with the connection was created
    /// with a [`config::Builder`].
    unsafe fn drop_config(&mut self) -> Result<(), Error> {
        let mut prev_config = core::ptr::null_mut();

        // A valid non-null pointer is returned only if the application previously called
        // [`Self::set_config()`].
        if s2n_connection_get_config(self.connection.as_ptr(), &mut prev_config)
            .into_result()
            .is_ok()
        {
            let prev_config = NonNull::new(prev_config).expect(
                "config should exist since the call to s2n_connection_get_config was successful",
            );
            drop(Config::from_raw(prev_config));
        }

        Ok(())
    }

    /// Associates a configuration object with a connection.
    pub fn set_config(&mut self, mut config: Config) -> Result<&mut Self, Error> {
        unsafe {
            // attempt to drop the currently set config
            self.drop_config()?;

            s2n_connection_set_config(self.connection.as_ptr(), config.as_mut_ptr())
                .into_result()?;

            debug_assert! {
                s2n_connection_get_config(self.connection.as_ptr(), &mut core::ptr::null_mut()).into_result().is_ok(),
                "s2n_connection_set_config was successful"
            };

            // Setting the config on the connection creates one additional reference to the config
            // so do not drop so prevent Rust from calling `drop()` at the end of this function.
            mem::forget(config);
        }

        Ok(self)
    }

    pub(crate) fn config(&self) -> Option<Config> {
        let mut raw = core::ptr::null_mut();
        let config = unsafe {
            s2n_connection_get_config(self.connection.as_ptr(), &mut raw)
                .into_result()
                .ok()?;
            let raw = NonNull::new(raw)?;
            Config::from_raw(raw)
        };
        // Because the config pointer is still set on the connection, this is a copy,
        // not the original config. This is fine -- Configs are immutable.
        let _ = ManuallyDrop::new(config.clone());
        Some(config)
    }

    pub fn set_security_policy(&mut self, policy: &security::Policy) -> Result<&mut Self, Error> {
        unsafe {
            s2n_connection_set_cipher_preferences(
                self.connection.as_ptr(),
                policy.as_cstr().as_ptr(),
            )
            .into_result()
        }?;
        Ok(self)
    }

    /// provides a smooth transition from s2n_connection_prefer_low_latency to s2n_connection_prefer_throughput.
    ///
    /// s2n_send uses small TLS records that fit into a single TCP segment for the resize_threshold
    /// bytes (cap to 8M) of data and reset record size back to a single segment after timeout_threshold
    /// seconds of inactivity.
    pub fn set_dynamic_record_threshold(
        &mut self,
        resize_threshold: u32,
        timeout_threshold: u16,
    ) -> Result<&mut Self, Error> {
        unsafe {
            s2n_connection_set_dynamic_record_threshold(
                self.connection.as_ptr(),
                resize_threshold,
                timeout_threshold,
            )
            .into_result()
        }?;
        Ok(self)
    }

    /// sets the application protocol preferences on an s2n_connection object.
    ///
    /// protocols is a list in order of preference, with most preferred protocol first, and of
    /// length protocol_count. When acting as a client the protocol list is included in the
    /// Client Hello message as the ALPN extension. As a server, the list is used to negotiate
    /// a mutual application protocol with the client. After the negotiation for the connection has
    /// completed, the agreed upon protocol can be retrieved with s2n_get_application_protocol
    pub fn set_application_protocol_preference<P: IntoIterator<Item = I>, I: AsRef<[u8]>>(
        &mut self,
        protocols: P,
    ) -> Result<&mut Self, Error> {
        // reset the list
        unsafe {
            s2n_connection_set_protocol_preferences(self.connection.as_ptr(), core::ptr::null(), 0)
                .into_result()
        }?;

        for protocol in protocols {
            self.append_application_protocol_preference(protocol.as_ref())?;
        }

        Ok(self)
    }

    pub fn append_application_protocol_preference(
        &mut self,
        protocol: &[u8],
    ) -> Result<&mut Self, Error> {
        unsafe {
            s2n_connection_append_protocol_preference(
                self.connection.as_ptr(),
                protocol.as_ptr(),
                protocol
                    .len()
                    .try_into()
                    .map_err(|_| Error::INVALID_INPUT)?,
            )
            .into_result()
        }?;
        Ok(self)
    }

    /// may be used to receive data with callbacks defined by the user.
    pub fn set_receive_callback(&mut self, callback: s2n_recv_fn) -> Result<&mut Self, Error> {
        unsafe { s2n_connection_set_recv_cb(self.connection.as_ptr(), callback).into_result() }?;
        Ok(self)
    }

    /// # Safety
    ///
    /// The `context` pointer must live at least as long as the connection
    pub unsafe fn set_receive_context(&mut self, context: *mut c_void) -> Result<&mut Self, Error> {
        s2n_connection_set_recv_ctx(self.connection.as_ptr(), context).into_result()?;
        Ok(self)
    }

    /// may be used to receive data with callbacks defined by the user.
    pub fn set_send_callback(&mut self, callback: s2n_send_fn) -> Result<&mut Self, Error> {
        unsafe { s2n_connection_set_send_cb(self.connection.as_ptr(), callback).into_result() }?;
        Ok(self)
    }

    /// Sets the callback to use for verifying that a hostname from an X.509 certificate is
    /// trusted.
    ///
    /// The callback may be called more than once during certificate validation as each SAN on
    /// the certificate will be checked.
    ///
    /// Corresponds to the underlying C API
    /// [s2n_connection_set_verify_host_callback](https://aws.github.io/s2n-tls/doxygen/s2n_8h.html).
    pub fn set_verify_host_callback<T: 'static + VerifyHostNameCallback>(
        &mut self,
        handler: T,
    ) -> Result<&mut Self, Error> {
        unsafe extern "C" fn verify_host_cb_fn(
            host_name: *const ::libc::c_char,
            host_name_len: usize,
            context: *mut ::libc::c_void,
        ) -> u8 {
            let context = &mut *(context as *mut Context);
            let handler = context.verify_host_callback.as_mut().unwrap();
            verify_host(host_name, host_name_len, handler)
        }

        self.context_mut().verify_host_callback = Some(Box::new(handler));
        unsafe {
            s2n_connection_set_verify_host_callback(
                self.connection.as_ptr(),
                Some(verify_host_cb_fn),
                self.context_mut() as *mut Context as *mut c_void,
            )
            .into_result()
        }?;
        Ok(self)
    }

    /// # Safety
    ///
    /// The `context` pointer must live at least as long as the connection
    pub unsafe fn set_send_context(&mut self, context: *mut c_void) -> Result<&mut Self, Error> {
        s2n_connection_set_send_ctx(self.connection.as_ptr(), context).into_result()?;
        Ok(self)
    }

    /// Connections prefering low latency will be encrypted using small record sizes that
    /// can be decrypted sooner by the recipient.
    pub fn prefer_low_latency(&mut self) -> Result<&mut Self, Error> {
        unsafe { s2n_connection_prefer_low_latency(self.connection.as_ptr()).into_result() }?;
        Ok(self)
    }

    /// Connections prefering throughput will use large record sizes that minimize overhead.
    pub fn prefer_throughput(&mut self) -> Result<&mut Self, Error> {
        unsafe { s2n_connection_prefer_throughput(self.connection.as_ptr()).into_result() }?;
        Ok(self)
    }

    /// wipes and free the in and out buffers associated with a connection.
    ///
    /// This function may be called when a connection is in keep-alive or idle state to
    /// reduce memory overhead of long lived connections.
    pub fn release_buffers(&mut self) -> Result<&mut Self, Error> {
        unsafe { s2n_connection_release_buffers(self.connection.as_ptr()).into_result() }?;
        Ok(self)
    }

    pub fn use_corked_io(&mut self) -> Result<&mut Self, Error> {
        unsafe { s2n_connection_use_corked_io(self.connection.as_ptr()).into_result() }?;
        Ok(self)
    }

    /// wipes an existing connection and allows it to be reused.
    ///
    /// This method erases all data associated with a connection including pending reads.
    /// This function should be called after all I/O is completed and s2n_shutdown has been
    /// called. Reusing the same connection handle(s) is more performant than repeatedly
    /// calling s2n_connection_new and s2n_connection_free
    pub fn wipe(&mut self) -> Result<&mut Self, Error> {
        let mode = self.mode();
        unsafe {
            // Wiping the connection will wipe the pointer to the context,
            // so retrieve and drop that memory first.
            let ctx = self.context_mut();
            drop(Box::from_raw(ctx));

            s2n_connection_wipe(self.connection.as_ptr()).into_result()
        }?;

        self.init_context(mode);
        Ok(self)
    }

    /// Performs the TLS handshake to completion
    ///
    /// Multiple callbacks can be configured for a connection and config, but
    /// [`Self::poll_negotiate()`] can only execute and block on one callback at a time.
    /// The handshake is sequential, not concurrent, and stops execution when
    /// it encounters an async callback.
    ///
    /// The handshake does not continue execution (and therefore can't call
    /// any other callbacks) until the blocking async task reports completion.
    pub fn poll_negotiate(&mut self) -> Poll<Result<&mut Self, Error>> {
        let mut blocked = s2n_blocked_status::NOT_BLOCKED;

        loop {
            // check if an async task exists and poll it to completion
            if let Some(fut) = self.poll_async_task() {
                match fut {
                    Poll::Ready(Ok(())) => {
                        // happy case:
                        // continue and call s2n_negotiate to make progress on the handshake
                    }
                    Poll::Ready(Err(err)) => {
                        // error case:
                        // if the callback returned an error then abort the handshake
                        return Poll::Ready(Err(err));
                    }
                    Poll::Pending => return Poll::Pending,
                }
            }

            let res = unsafe { s2n_negotiate(self.connection.as_ptr(), &mut blocked).into_poll() };

            match res {
                Poll::Ready(res) => {
                    let res = res.map(|_| self);
                    return Poll::Ready(res);
                }
                Poll::Pending => {
                    // if there is no connection_future then return, otherwise continue
                    // looping and polling the future
                    if self.context_mut().async_callback.is_none() {
                        return Poll::Pending;
                    }
                }
            }
        }
    }

    // Poll the connection future if it exists.
    //
    // If the future returns Pending, then re-set it back on the Connection.
    fn poll_async_task(&mut self) -> Option<Poll<Result<(), Error>>> {
        self.take_async_callback().map(|mut callback| {
            let waker = self.waker().ok_or(Error::MISSING_WAKER)?.clone();
            let mut ctx = core::task::Context::from_waker(&waker);
            match Pin::new(&mut callback).poll(self, &mut ctx) {
                Poll::Ready(result) => Poll::Ready(result),
                Poll::Pending => {
                    // replace the future if it hasn't completed yet
                    self.set_async_callback(callback);
                    Poll::Pending
                }
            }
        })
    }

    /// Encrypts and sends data on a connection where
    /// [negotiate](`Self::poll_negotiate`) has succeeded.
    ///
    /// Returns the number of bytes written, and may indicate a partial write.
    pub fn poll_send(&mut self, buf: &[u8]) -> Poll<Result<usize, Error>> {
        let mut blocked = s2n_blocked_status::NOT_BLOCKED;
        let buf_len: isize = buf.len().try_into().map_err(|_| Error::INVALID_INPUT)?;
        let buf_ptr = buf.as_ptr() as *const ::libc::c_void;
        unsafe { s2n_send(self.connection.as_ptr(), buf_ptr, buf_len, &mut blocked).into_poll() }
    }

    /// Reads and decrypts data from a connection where
    /// [negotiate](`Self::poll_negotiate`) has succeeded.
    ///
    /// Returns the number of bytes read, and may indicate a partial read.
    /// 0 bytes returned indicates EOF due to connection closure.
    pub fn poll_recv(&mut self, buf: &mut [u8]) -> Poll<Result<usize, Error>> {
        let mut blocked = s2n_blocked_status::NOT_BLOCKED;
        let buf_len: isize = buf.len().try_into().map_err(|_| Error::INVALID_INPUT)?;
        let buf_ptr = buf.as_ptr() as *mut ::libc::c_void;
        unsafe { s2n_recv(self.connection.as_ptr(), buf_ptr, buf_len, &mut blocked).into_poll() }
    }

    /// Reads and decrypts data from a connection where
    /// [negotiate](`Self::poll_negotiate`) has succeeded
    /// to a uninitialized buffer.
    ///
    /// Returns the number of bytes read, and may indicate a partial read.
    /// 0 bytes returned indicates EOF due to connection closure.
    ///
    /// Safety: this function is always safe to call, and additionally:
    /// 1. It will never deinitialize any bytes in `buf`.
    /// 2. If it returns `Ok(n)`, then the first `n` bytes of `buf`
    /// will have been initialized by this function.
    pub fn poll_recv_uninitialized(
        &mut self,
        buf: &mut [MaybeUninit<u8>],
    ) -> Poll<Result<usize, Error>> {
        let mut blocked = s2n_blocked_status::NOT_BLOCKED;
        let buf_len: isize = buf.len().try_into().map_err(|_| Error::INVALID_INPUT)?;
        let buf_ptr = buf.as_ptr() as *mut ::libc::c_void;

        // Safety:
        // 1. s2n_recv never writes uninitialized garbage to `buf`.
        // 2. if s2n_recv returns `+n`, it guarantees that the first
        // `n` bytes of `buf` have been initialized, which allows this
        // function to return `Ok(n)`
        unsafe { s2n_recv(self.connection.as_ptr(), buf_ptr, buf_len, &mut blocked).into_poll() }
    }

    /// Attempts to flush any data previously buffered by a call to [send](`Self::poll_send`).
    pub fn poll_flush(&mut self) -> Poll<Result<&mut Self, Error>> {
        self.poll_send(&[0; 0]).map_ok(|_| self)
    }

    /// Gets the number of bytes that are currently available in the buffer to be read.
    pub fn peek_len(&self) -> usize {
        unsafe { s2n_peek(self.connection.as_ptr()) as usize }
    }

    /// Attempts a graceful shutdown of the TLS connection.
    ///
    /// The shutdown is not complete until the necessary shutdown messages
    /// have been successfully sent and received. If the peer does not respond
    /// correctly, the graceful shutdown may fail.
    pub fn poll_shutdown(&mut self) -> Poll<Result<&mut Self, Error>> {
        if !self.remaining_blinding_delay()?.is_zero() {
            return Poll::Pending;
        }
        let mut blocked = s2n_blocked_status::NOT_BLOCKED;
        unsafe {
            s2n_shutdown(self.connection.as_ptr(), &mut blocked)
                .into_poll()
                .map_ok(|_| self)
        }
    }

    /// Returns the TLS alert code, if any
    pub fn alert(&self) -> Option<u8> {
        let alert =
            unsafe { s2n_connection_get_alert(self.connection.as_ptr()).into_result() }.ok()?;
        Some(alert as u8)
    }

    /// Sets the server name value for the connection
    pub fn set_server_name(&mut self, server_name: &str) -> Result<&mut Self, Error> {
        let server_name = std::ffi::CString::new(server_name).map_err(|_| Error::INVALID_INPUT)?;
        unsafe {
            s2n_set_server_name(self.connection.as_ptr(), server_name.as_ptr()).into_result()
        }?;
        Ok(self)
    }

    /// Get the server name associated with the connection client hello.
    pub fn server_name(&self) -> Option<&str> {
        unsafe {
            let server_name = s2n_get_server_name(self.connection.as_ptr());
            match server_name.into_result() {
                Ok(server_name) => CStr::from_ptr(server_name).to_str().ok(),
                Err(_) => None,
            }
        }
    }

    /// Sets a Waker on the connection context or clears it if `None` is passed.
    pub fn set_waker(&mut self, waker: Option<&Waker>) -> Result<&mut Self, Error> {
        let ctx = self.context_mut();

        if let Some(waker) = waker {
            if let Some(prev_waker) = ctx.waker.as_mut() {
                // only replace the Waker if they dont reference the same task
                if !prev_waker.will_wake(waker) {
                    *prev_waker = waker.clone();
                }
            } else {
                ctx.waker = Some(waker.clone());
            }
        } else {
            ctx.waker = None;
        }
        Ok(self)
    }

    /// Returns the Waker set on the connection context.
    pub fn waker(&self) -> Option<&Waker> {
        let ctx = self.context();
        ctx.waker.as_ref()
    }

    /// Takes the [`Option::take`] the connection_future stored on the
    /// connection context.
    ///
    /// If the Future returns `Poll::Pending` and has not completed, then it
    /// should be re-set using [`Self::set_connection_future()`]
    fn take_async_callback(&mut self) -> Option<AsyncCallback> {
        let ctx = self.context_mut();
        ctx.async_callback.take()
    }

    /// Sets a `connection_future` on the connection context.
    pub(crate) fn set_async_callback(&mut self, callback: AsyncCallback) {
        let ctx = self.context_mut();
        debug_assert!(ctx.async_callback.is_none());
        ctx.async_callback = Some(callback);
    }

    /// Retrieve a mutable reference to the [`Context`] stored on the connection.
    fn context_mut(&mut self) -> &mut Context {
        unsafe {
            let ctx = s2n_connection_get_ctx(self.connection.as_ptr())
                .into_result()
                .unwrap();
            &mut *(ctx.as_ptr() as *mut Context)
        }
    }

    /// Retrieve a reference to the [`Context`] stored on the connection.
    fn context(&self) -> &Context {
        unsafe {
            let ctx = s2n_connection_get_ctx(self.connection.as_ptr())
                .into_result()
                .unwrap();
            &*(ctx.as_ptr() as *mut Context)
        }
    }

    /// Mark that the server_name extension was used to configure the connection.
    pub fn server_name_extension_used(&mut self) {
        // TODO: requiring the application to call this method is a pretty sharp edge.
        // Figure out if its possible to automatically call this from the Rust bindings.
        unsafe {
            s2n_connection_server_name_extension_used(self.connection.as_ptr())
                .into_result()
                .unwrap();
        }
    }

    /// Check if client auth was used for a connection.
    ///
    /// This is only relevant if [`ClientAuthType::Optional] was used.
    pub fn client_cert_used(&self) -> bool {
        unsafe { s2n_connection_client_cert_used(self.connection.as_ptr()) == 1 }
    }

    /// Retrieves the raw bytes of the client cert chain received from the peer, if present.
    pub fn client_cert_chain_bytes(&self) -> Result<Option<&[u8]>, Error> {
        if !self.client_cert_used() {
            return Ok(None);
        }

        let mut chain = std::ptr::null_mut();
        let mut len = 0;
        unsafe {
            s2n_connection_get_client_cert_chain(self.connection.as_ptr(), &mut chain, &mut len)
                .into_result()?;
        }

        if chain.is_null() || len == 0 {
            return Ok(None);
        }

        unsafe { Ok(Some(std::slice::from_raw_parts(chain, len as usize))) }
    }

    // The memory backing the ClientHello is owned by the Connection, so we
    // tie the ClientHello to the lifetime of the Connection. This is validated
    // with a doc test that ensures the ClientHello is invalid once the
    // connection has gone out of scope.
    //
    /// Returns a reference to the ClientHello associated with the connection.
    /// ```compile_fail
    /// use s2n_tls::client_hello::{ClientHello, FingerprintType};
    /// use s2n_tls::connection::Connection;
    /// use s2n_tls::enums::Mode;
    ///
    /// let mut conn = Connection::new(Mode::Server);
    /// let mut client_hello: &ClientHello = conn.client_hello().unwrap();
    /// let mut hash = Vec::new();
    /// drop(conn);
    /// client_hello.fingerprint_hash(FingerprintType::JA3, &mut hash);
    /// ```
    ///
    /// The compilation could be failing for a variety of reasons, so make sure
    /// that the test case is actually good.
    /// ```no_run
    /// use s2n_tls::client_hello::{ClientHello, FingerprintType};
    /// use s2n_tls::connection::Connection;
    /// use s2n_tls::enums::Mode;
    ///
    /// let mut conn = Connection::new(Mode::Server);
    /// let mut client_hello: &ClientHello = conn.client_hello().unwrap();
    /// let mut hash = Vec::new();
    /// client_hello.fingerprint_hash(FingerprintType::JA3, &mut hash);
    /// drop(conn);
    /// ```
    #[cfg(feature = "unstable-fingerprint")]
    pub fn client_hello(&self) -> Result<&crate::client_hello::ClientHello, Error> {
        let mut handle =
            unsafe { s2n_connection_get_client_hello(self.connection.as_ptr()).into_result()? };
        Ok(crate::client_hello::ClientHello::from_ptr(unsafe {
            handle.as_mut()
        }))
    }

    pub(crate) fn mark_client_hello_cb_done(&mut self) -> Result<(), Error> {
        unsafe {
            s2n_client_hello_cb_done(self.connection.as_ptr()).into_result()?;
        }
        Ok(())
    }

    pub fn actual_protocol_version(&self) -> Result<Version, Error> {
        let version = unsafe {
            s2n_connection_get_actual_protocol_version(self.connection.as_ptr()).into_result()?
        };
        version.try_into()
    }

    pub fn handshake_type(&self) -> Result<&str, Error> {
        let handshake = unsafe {
            s2n_connection_get_handshake_type_name(self.connection.as_ptr()).into_result()?
        };
        // The strings returned by s2n_connection_get_handshake_type_name
        // are static and immutable after they are first calculated
        static_const_str!(handshake)
    }

    pub fn cipher_suite(&self) -> Result<&str, Error> {
        let cipher = unsafe { s2n_connection_get_cipher(self.connection.as_ptr()).into_result()? };
        // The strings returned by s2n_connection_get_cipher
        // are static and immutable since they are const fields on static const structs
        static_const_str!(cipher)
    }

    pub fn selected_curve(&self) -> Result<&str, Error> {
        let curve = unsafe { s2n_connection_get_curve(self.connection.as_ptr()).into_result()? };
        static_const_str!(curve)
    }

    pub fn selected_signature_algorithm(&self) -> Result<SignatureAlgorithm, Error> {
        let mut sig_alg = s2n_tls_signature_algorithm::ANONYMOUS;
        unsafe {
            s2n_connection_get_selected_signature_algorithm(self.connection.as_ptr(), &mut sig_alg)
                .into_result()?;
        }
        sig_alg.try_into()
    }

    pub fn selected_hash_algorithm(&self) -> Result<HashAlgorithm, Error> {
        let mut hash_alg = s2n_tls_hash_algorithm::NONE;
        unsafe {
            s2n_connection_get_selected_digest_algorithm(self.connection.as_ptr(), &mut hash_alg)
                .into_result()?;
        }
        hash_alg.try_into()
    }

    pub fn selected_client_signature_algorithm(&self) -> Result<Option<SignatureAlgorithm>, Error> {
        let mut sig_alg = s2n_tls_signature_algorithm::ANONYMOUS;
        unsafe {
            s2n_connection_get_selected_client_cert_signature_algorithm(
                self.connection.as_ptr(),
                &mut sig_alg,
            )
            .into_result()?;
        }
        Ok(match sig_alg {
            s2n_tls_signature_algorithm::ANONYMOUS => None,
            sig_alg => Some(sig_alg.try_into()?),
        })
    }

    pub fn selected_client_hash_algorithm(&self) -> Result<Option<HashAlgorithm>, Error> {
        let mut hash_alg = s2n_tls_hash_algorithm::NONE;
        unsafe {
            s2n_connection_get_selected_client_cert_digest_algorithm(
                self.connection.as_ptr(),
                &mut hash_alg,
            )
            .into_result()?;
        }
        Ok(match hash_alg {
            s2n_tls_hash_algorithm::NONE => None,
            hash_alg => Some(hash_alg.try_into()?),
        })
    }
}

struct Context {
    mode: Mode,
    waker: Option<Waker>,
    async_callback: Option<AsyncCallback>,
    verify_host_callback: Option<Box<dyn VerifyHostNameCallback>>,
}

impl Context {
    fn new(mode: Mode) -> Self {
        Context {
            mode,
            waker: None,
            async_callback: None,
            verify_host_callback: None,
        }
    }
}

#[cfg(feature = "quic")]
impl Connection {
    pub fn enable_quic(&mut self) -> Result<&mut Self, Error> {
        unsafe { s2n_connection_enable_quic(self.connection.as_ptr()).into_result() }?;
        Ok(self)
    }

    pub fn set_quic_transport_parameters(&mut self, buffer: &[u8]) -> Result<&mut Self, Error> {
        unsafe {
            s2n_connection_set_quic_transport_parameters(
                self.connection.as_ptr(),
                buffer.as_ptr(),
                buffer.len().try_into().map_err(|_| Error::INVALID_INPUT)?,
            )
            .into_result()
        }?;
        Ok(self)
    }

    pub fn quic_transport_parameters(&mut self) -> Result<&[u8], Error> {
        let mut ptr = core::ptr::null();
        let mut len = 0;
        unsafe {
            s2n_connection_get_quic_transport_parameters(
                self.connection.as_ptr(),
                &mut ptr,
                &mut len,
            )
            .into_result()
        }?;
        let buffer = unsafe { core::slice::from_raw_parts(ptr, len as _) };
        Ok(buffer)
    }

    /// # Safety
    ///
    /// The `context` pointer must live at least as long as the connection
    pub unsafe fn set_secret_callback(
        &mut self,
        callback: s2n_secret_cb,
        context: *mut c_void,
    ) -> Result<&mut Self, Error> {
        s2n_connection_set_secret_callback(self.connection.as_ptr(), callback, context)
            .into_result()?;
        Ok(self)
    }
}

impl AsRef<Connection> for Connection {
    fn as_ref(&self) -> &Connection {
        self
    }
}

impl AsMut<Connection> for Connection {
    fn as_mut(&mut self) -> &mut Connection {
        self
    }
}

impl Drop for Connection {
    fn drop(&mut self) {
        // ignore failures since there's not much we can do about it
        unsafe {
            // clean up context
            let prev_ctx = self.context_mut();
            drop(Box::from_raw(prev_ctx));
            let _ = s2n_connection_set_ctx(self.connection.as_ptr(), core::ptr::null_mut())
                .into_result();

            // cleanup config
            let _ = self.drop_config();

            // cleanup connection
            let _ = s2n_connection_free(self.connection.as_ptr()).into_result();
        }
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

#![allow(clippy::missing_safety_doc)] // TODO add safety docs

use crate::error::Error;
use core::convert::TryFrom;
use s2n_tls_sys::*;

#[derive(Debug, PartialEq, Copy, Clone)]
pub enum CallbackResult {
    Success,
    Failure,
}

impl From<CallbackResult> for s2n_status_code::Type {
    fn from(input: CallbackResult) -> s2n_status_code::Type {
        match input {
            CallbackResult::Success => s2n_status_code::SUCCESS,
            CallbackResult::Failure => s2n_status_code::FAILURE,
        }
    }
}

impl<T, E> From<Result<T, E>> for CallbackResult {
    fn from(result: Result<T, E>) -> CallbackResult {
        match result {
            Ok(_) => CallbackResult::Success,
            Err(_) => CallbackResult::Failure,
        }
    }
}

#[derive(Debug, PartialEq, Copy, Clone)]
pub enum Mode {
    Server,
    Client,
}

impl From<Mode> for s2n_mode::Type {
    fn from(input: Mode) -> s2n_mode::Type {
        match input {
            Mode::Server => s2n_mode::SERVER,
            Mode::Client => s2n_mode::CLIENT,
        }
    }
}

#[non_exhaustive]
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum Version {
    SSLV2,
    SSLV3,
    TLS10,
    TLS11,
    TLS12,
    TLS13,
}

impl TryFrom<s2n_tls_version::Type> for Version {
    type Error = Error;

    fn try_from(input: s2n_tls_version::Type) -> Result<Self, Self::Error> {
        let version = match input {
            s2n_tls_version::SSLV2 => Self::SSLV2,
            s2n_tls_version::SSLV3 => Self::SSLV3,
            s2n_tls_version::TLS10 => Self::TLS10,
            s2n_tls_version::TLS11 => Self::TLS11,
            s2n_tls_version::TLS12 => Self::TLS12,
            s2n_tls_version::TLS13 => Self::TLS13,
            _ => return Err(Error::INVALID_INPUT),
        };
        Ok(version)
    }
}

#[non_exhaustive]
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum Blinding {
    SelfService,
    BuiltIn,
}

impl From<Blinding> for s2n_blinding::Type {
    fn from(input: Blinding) -> s2n_blinding::Type {
        match input {
            Blinding::SelfService => s2n_blinding::SELF_SERVICE_BLINDING,
            Blinding::BuiltIn => s2n_blinding::BUILT_IN_BLINDING,
        }
    }
}

#[non_exhaustive]
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum ClientAuthType {
    Required,
    Optional,
    None,
}

impl From<ClientAuthType> for s2n_cert_auth_type::Type {
    fn from(input: ClientAuthType) -> s2n_cert_auth_type::Type {
        match input {
            ClientAuthType::Required => s2n_cert_auth_type::REQUIRED,
            ClientAuthType::Optional => s2n_cert_auth_type::OPTIONAL,
            ClientAuthType::None => s2n_cert_auth_type::NONE,
        }
    }
}

#[non_exhaustive]
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum AlertBehavior {
    FailOnWarnings,
    IgnoreWarnings,
}

impl From<AlertBehavior> for s2n_alert_behavior::Type {
    fn from(input: AlertBehavior) -> s2n_alert_behavior::Type {
        match input {
            AlertBehavior::FailOnWarnings => s2n_alert_behavior::FAIL_ON_WARNINGS,
            AlertBehavior::IgnoreWarnings => s2n_alert_behavior::IGNORE_WARNINGS,
        }
    }
}

#[non_exhaustive]
#[derive(Debug, PartialEq, Copy, Clone)]
#[allow(non_camel_case_types)]
pub enum SignatureAlgorithm {
    RSA_PKCS1,
    RSA_PSS_RSAE,
    RSA_PSS_PSS,
    ECDSA,
}

impl TryFrom<s2n_tls_signature_algorithm::Type> for SignatureAlgorithm {
    type Error = Error;

    fn try_from(input: s2n_tls_signature_algorithm::Type) -> Result<Self, Self::Error> {
        let version = match input {
            s2n_tls_signature_algorithm::RSA => Self::RSA_PKCS1,
            s2n_tls_signature_algorithm::RSA_PSS_RSAE => Self::RSA_PSS_RSAE,
            s2n_tls_signature_algorithm::RSA_PSS_PSS => Self::RSA_PSS_PSS,
            s2n_tls_signature_algorithm::ECDSA => Self::ECDSA,
            _ => return Err(Error::INVALID_INPUT),
        };
        Ok(version)
    }
}

#[non_exhaustive]
#[derive(Debug, PartialEq, Copy, Clone)]
#[allow(non_camel_case_types)]
pub enum HashAlgorithm {
    MD5,
    SHA1,
    SHA224,
    SHA256,
    SHA384,
    SHA512,
}

impl TryFrom<s2n_tls_hash_algorithm::Type> for HashAlgorithm {
    type Error = Error;

    fn try_from(input: s2n_tls_hash_algorithm::Type) -> Result<Self, Self::Error> {
        let version = match input {
            s2n_tls_hash_algorithm::MD5 => Self::MD5,
            s2n_tls_hash_algorithm::SHA1 => Self::SHA1,
            s2n_tls_hash_algorithm::SHA224 => Self::SHA224,
            s2n_tls_hash_algorithm::SHA256 => Self::SHA256,
            s2n_tls_hash_algorithm::SHA384 => Self::SHA384,
            s2n_tls_hash_algorithm::SHA512 => Self::SHA512,
            _ => return Err(Error::INVALID_INPUT),
        };
        Ok(version)
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use core::{convert::TryInto, fmt, ptr::NonNull, task::Poll};
use errno::{errno, Errno};
use libc::c_char;
use s2n_tls_sys::*;
use std::{convert::TryFrom, ffi::CStr};

#[non_exhaustive]
#[derive(Debug, PartialEq)]
pub enum ErrorType {
    UnknownErrorType,
    NoError,
    IOError,
    ConnectionClosed,
    Blocked,
    Alert,
    ProtocolError,
    InternalError,
    UsageError,
    Application,
}

#[non_exhaustive]
#[derive(Debug, PartialEq)]
pub enum ErrorSource {
    Library,
    Bindings,
    Application,
}

impl From<libc::c_int> for ErrorType {
    fn from(input: libc::c_int) -> Self {
        match input as s2n_error_type::Type {
            s2n_error_type::OK => ErrorType::NoError,
            s2n_error_type::IO => ErrorType::IOError,
            s2n_error_type::CLOSED => ErrorType::ConnectionClosed,
            s2n_error_type::BLOCKED => ErrorType::Blocked,
            s2n_error_type::ALERT => ErrorType::Alert,
            s2n_error_type::PROTO => ErrorType::ProtocolError,
            s2n_error_type::INTERNAL => ErrorType::InternalError,
            s2n_error_type::USAGE => ErrorType::UsageError,
            _ => ErrorType::UnknownErrorType,
        }
    }
}

enum Context {
    InvalidInput,
    MissingWaker,
    Code(s2n_status_code::Type, Errno),
    Application(Box<dyn std::error::Error + Send + Sync + 'static>),
}

pub struct Error(Context);

pub trait Fallible {
    type Output;

    fn into_result(self) -> Result<Self::Output, Error>;
}

impl Fallible for s2n_status_code::Type {
    type Output = s2n_status_code::Type;

    fn into_result(self) -> Result<Self::Output, Error> {
        if self >= s2n_status_code::SUCCESS {
            Ok(self)
        } else {
            Err(Error::capture())
        }
    }
}

impl Fallible for isize {
    type Output = usize;

    fn into_result(self) -> Result<Self::Output, Error> {
        // Negative values can't be converted to a real size
        // and instead indicate an error.
        self.try_into().map_err(|_| Error::capture())
    }
}

impl Fallible for u64 {
    type Output = Self;

    /// Converts a u64 to a Result by checking for u64::MAX.
    ///
    /// If a method that returns an unsigned int is fallible,
    /// then the -1 error result wraps around to u64::MAX.
    ///
    /// For a u64 to be Fallible, a result of u64::MAX must not be
    /// possible without an error. For example, [`s2n_connection_get_delay`]
    /// can't return u64::MAX as a valid result because
    /// s2n-tls blinding delays are limited to 30s, or a return value of 3^10 ns,
    /// which is significantly less than u64::MAX. [`s2n_connection_get_delay`]
    /// would therefore only return u64::MAX for a -1 error result.
    fn into_result(self) -> Result<Self::Output, Error> {
        if self != Self::MAX {
            Ok(self)
        } else {
            Err(Error::capture())
        }
    }
}

impl<T> Fallible for *mut T {
    type Output = NonNull<T>;

    fn into_result(self) -> Result<Self::Output, Error> {
        if let Some(value) = NonNull::new(self) {
            Ok(value)
        } else {
            Err(Error::capture())
        }
    }
}

impl<T> Fallible for *const T {
    type Output = *const T;

    fn into_result(self) -> Result<Self::Output, Error> {
        if !self.is_null() {
            Ok(self)
        } else {
            Err(Error::capture())
        }
    }
}

pub trait Pollable {
    type Output;

    fn into_poll(self) -> Poll<Result<Self::Output, Error>>;
}

impl<T: Fallible> Pollable for T {
    type Output = T::Output;

    fn into_poll(self) -> Poll<Result<Self::Output, Error>> {
        match self.into_result() {
            Ok(r) => Ok(r).into(),
            Err(err) if err.is_retryable() => Poll::Pending,
            Err(err) => Err(err).into(),
        }
    }
}

impl Error {
    pub(crate) const INVALID_INPUT: Error = Self(Context::InvalidInput);
    pub(crate) const MISSING_WAKER: Error = Self(Context::MissingWaker);

    /// Converts an io::Error into an s2n-tls Error
    pub fn io_error(err: std::io::Error) -> Error {
        let errno = err.raw_os_error().unwrap_or(1);
        errno::set_errno(errno::Errno(errno));
        s2n_status_code::FAILURE.into_result().unwrap_err()
    }

    /// An error occurred while running application code.
    ///
    /// Can be emitted from [`crate::callbacks::ConnectionFuture::poll()`] to indicate
    /// async task failure.
    pub fn application(error: Box<dyn std::error::Error + Send + Sync + 'static>) -> Self {
        Self(Context::Application(error))
    }

    fn capture() -> Self {
        unsafe {
            let s2n_errno = s2n_errno_location();

            let code = *s2n_errno;

            // https://github.com/aws/s2n-tls/blob/main/docs/USAGE-GUIDE.md#error-handling
            //# To avoid possible confusion, s2n_errno should be cleared after processing
            //# an error: s2n_errno = S2N_ERR_T_OK
            *s2n_errno = s2n_error_type::OK as _;

            Self(Context::Code(code, errno()))
        }
    }

    pub fn name(&self) -> &'static str {
        match self.0 {
            Context::InvalidInput => "InvalidInput",
            Context::MissingWaker => "MissingWaker",
            Context::Application(_) => "ApplicationError",
            Context::Code(code, _) => unsafe {
                // Safety: we assume the string has a valid encoding coming from s2n
                cstr_to_str(s2n_strerror_name(code))
            },
        }
    }

    pub fn message(&self) -> &'static str {
        match self.0 {
            Context::InvalidInput => "A parameter was incorrect",
            Context::MissingWaker => {
                "Tried to perform an asynchronous operation without a configured waker"
            }
            Context::Application(_) => "An error occurred while executing application code",
            Context::Code(code, _) => unsafe {
                // Safety: we assume the string has a valid encoding coming from s2n
                cstr_to_str(s2n_strerror(code, core::ptr::null()))
            },
        }
    }

    pub fn debug(&self) -> Option<&'static str> {
        match self.0 {
            Context::InvalidInput | Context::MissingWaker | Context::Application(_) => None,
            Context::Code(code, _) => unsafe {
                let debug_info = s2n_strerror_debug(code, core::ptr::null());

                // The debug string should be set to a constant static string
                // when an error occurs, but because it starts out as NULL
                // we should defend against mistakes.
                if debug_info.is_null() {
                    None
                } else {
                    // If the string is not null, then we can assume that
                    // it is constant and static.
                    Some(cstr_to_str(debug_info))
                }
            },
        }
    }

    pub fn kind(&self) -> ErrorType {
        match self.0 {
            Context::InvalidInput | Context::MissingWaker => ErrorType::UsageError,
            Context::Application(_) => ErrorType::Application,
            Context::Code(code, _) => unsafe { ErrorType::from(s2n_error_get_type(code)) },
        }
    }

    pub fn source(&self) -> ErrorSource {
        match self.0 {
            Context::InvalidInput | Context::MissingWaker => ErrorSource::Bindings,
            Context::Application(_) => ErrorSource::Application,
            Context::Code(_, _) => ErrorSource::Library,
        }
    }

    #[allow(clippy::borrowed_box)]
    /// Returns an [`std::error::Error`] if the error source was [`ErrorSource::Application`],
    /// otherwise returns None.
    pub fn application_error(&self) -> Option<&Box<dyn std::error::Error + Send + Sync + 'static>> {
        if let Self(Context::Application(err)) = self {
            Some(err)
        } else {
            None
        }
    }

    pub fn is_retryable(&self) -> bool {
        matches!(self.kind(), ErrorType::Blocked)
    }
}

#[cfg(feature = "quic")]
impl Error {
    /// s2n-tls does not send specific errors.
    ///
    /// However, we can attempt to map local errors into the alerts
    /// that we would have sent if we sent alerts.
    ///
    /// This API is currently incomplete and should not be relied upon.
    pub fn alert(&self) -> Option<u8> {
        match self.0 {
            Context::InvalidInput | Context::MissingWaker | Context::Application(_) => None,
            Context::Code(code, _) => {
                let mut alert = 0;
                let r = unsafe { s2n_error_get_alert(code, &mut alert) };
                match r.into_result() {
                    Ok(_) => Some(alert),
                    Err(_) => None,
                }
            }
        }
    }
}

/// # Safety
///
/// The caller must ensure the char pointer must contain a valid
/// UTF-8 string from a trusted source
unsafe fn cstr_to_str(v: *const c_char) -> &'static str {
    let slice = CStr::from_ptr(v);
    let bytes = slice.to_bytes();
    core::str::from_utf8_unchecked(bytes)
}

impl TryFrom<std::io::Error> for Error {
    type Error = Error;
    fn try_from(value: std::io::Error) -> Result<Self, Self::Error> {
        let io_inner = value.into_inner().ok_or(Error::INVALID_INPUT)?;
        io_inner
            .downcast::<Self>()
            .map(|error| *error)
            .map_err(|_| Error::INVALID_INPUT)
    }
}

impl From<Error> for std::io::Error {
    fn from(input: Error) -> Self {
        if let Context::Code(_, errno) = input.0 {
            if ErrorType::IOError == input.kind() {
                let bare = std::io::Error::from_raw_os_error(errno.0);
                return std::io::Error::new(bare.kind(), input);
            }
        }
        std::io::Error::new(std::io::ErrorKind::Other, input)
    }
}

impl fmt::Debug for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut s = f.debug_struct("Error");
        if let Context::Code(code, _) = self.0 {
            s.field("code", &code);
        }

        s.field("name", &self.name());
        s.field("message", &self.message());
        s.field("kind", &self.kind());
        s.field("source", &self.source());

        if let Some(debug) = self.debug() {
            s.field("debug", &debug);
        }

        // "errno" is only known to be meaningful for IOErrors.
        // However, it has occasionally proved useful for debugging
        // other errors, so include it for all errors.
        if let Context::Code(_, errno) = self.0 {
            s.field("errno", &errno.to_string());
        }

        s.finish()
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if let Self(Context::Application(err)) = self {
            err.fmt(f)
        } else {
            f.write_str(self.message())
        }
    }
}

impl std::error::Error for Error {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        // implement `source` in the same way `std::io::Error` implements it:
        // https://doc.rust-lang.org/std/io/struct.Error.html#method.source
        if let Self(Context::Application(err)) = self {
            err.source()
        } else {
            None
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{enums::Version, testing::client_hello::CustomError};
    use errno::set_errno;

    const FAILURE: isize = -1;

    // This relies on an implementation detail of s2n-tls errors,
    // and could make these tests brittle. However, the alternative
    // is a real handshake producing a real IO error, so just updating
    // this value if the definition of an IO error changes might be easier.
    const S2N_IO_ERROR_CODE: s2n_status_code::Type = 1 << 26;

    #[test]
    fn s2n_io_error_to_std_io_error() -> Result<(), Box<dyn std::error::Error>> {
        set_errno(Errno(libc::ECONNRESET));
        unsafe {
            let s2n_errno_ptr = s2n_errno_location();
            *s2n_errno_ptr = S2N_IO_ERROR_CODE;
        }

        let s2n_error = FAILURE.into_result().unwrap_err();
        assert_eq!(ErrorType::IOError, s2n_error.kind());

        let io_error = std::io::Error::from(s2n_error);
        assert_eq!(std::io::ErrorKind::ConnectionReset, io_error.kind());
        assert!(io_error.into_inner().is_some());
        Ok(())
    }

    #[test]
    fn s2n_error_to_std_io_error() -> Result<(), Box<dyn std::error::Error>> {
        set_errno(Errno(libc::ECONNRESET));
        unsafe {
            let s2n_errno_ptr = s2n_errno_location();
            *s2n_errno_ptr = S2N_IO_ERROR_CODE - 1;
        }

        let s2n_error = FAILURE.into_result().unwrap_err();
        assert_ne!(ErrorType::IOError, s2n_error.kind());

        let io_error = std::io::Error::from(s2n_error);
        assert_eq!(std::io::ErrorKind::Other, io_error.kind());
        assert!(io_error.into_inner().is_some());
        Ok(())
    }

    #[test]
    fn invalid_input_to_std_io_error() -> Result<(), Box<dyn std::error::Error>> {
        let s2n_error = Version::try_from(0).unwrap_err();
        assert_eq!(ErrorType::UsageError, s2n_error.kind());

        let io_error = std::io::Error::from(s2n_error);
        assert_eq!(std::io::ErrorKind::Other, io_error.kind());
        assert!(io_error.into_inner().is_some());
        Ok(())
    }

    #[test]
    fn error_source() -> Result<(), Box<dyn std::error::Error>> {
        let bindings_error = Version::try_from(0).unwrap_err();
        assert_eq!(ErrorSource::Bindings, bindings_error.source());

        let library_error = FAILURE.into_result().unwrap_err();
        assert_eq!(ErrorSource::Library, library_error.source());

        Ok(())
    }

    #[test]
    fn application_error() {
        // test single level errors
        {
            let error = Error::application(Box::new(CustomError));

            let app_error = error.application_error().unwrap();
            let _custom_error = app_error.downcast_ref::<CustomError>().unwrap();
        }

        // make sure nested errors work
        {
            let io_error = std::io::Error::new(std::io::ErrorKind::Other, CustomError);
            let error = Error::application(Box::new(io_error));

            let app_error = error.application_error().unwrap();
            let io_error = app_error.downcast_ref::<std::io::Error>().unwrap();
            let _custom_error = io_error
                .get_ref()
                .unwrap()
                .downcast_ref::<CustomError>()
                .unwrap();
        }
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::error::{Error, Fallible};
use s2n_tls_sys::*;
use std::sync::Once;

static S2N_INIT: Once = Once::new();

/// # Safety
///
/// This function should only be called once
unsafe fn global_init() -> Result<(), Error> {
    mem::init()?;
    s2n_init().into_result()?;
    Ok(())
}

thread_local! {
    static S2N_THREAD: Thread = {
        S2N_INIT.call_once(|| unsafe {
            // Safety: by using `Once` we can ensure the library is initialized once
            global_init().expect("could not initialize s2n-tls");
        });
        Thread
    };
}

struct Thread;

impl Drop for Thread {
    fn drop(&mut self) {
        // https://doc.rust-lang.org/std/thread/struct.LocalKey.html#platform-specific-behavior
        // Note that a "best effort" is made to ensure that destructors for types stored in thread local storage are run, but not all platforms can guarantee that destructors will be run for all types in thread local storage.
        let _ = unsafe { s2n_cleanup().into_result() };
    }
}

pub fn init() {
    S2N_THREAD.with(|_| ());
}

mod mem {
    use super::*;
    use alloc::alloc::{alloc, dealloc, Layout};
    use core::{ffi::c_void, mem::size_of};

    pub unsafe fn init() -> Result<(), Error> {
        s2n_mem_set_callbacks(
            Some(mem_init_callback),
            Some(mem_cleanup_callback),
            Some(mem_malloc_callback),
            Some(mem_free_callback),
        )
        .into_result()?;
        Ok(())
    }

    unsafe extern "C" fn mem_init_callback() -> s2n_status_code::Type {
        // no-op: the global allocator is already initialized
        s2n_status_code::SUCCESS
    }

    unsafe extern "C" fn mem_cleanup_callback() -> s2n_status_code::Type {
        // no-op: the global allocator is already initialized
        s2n_status_code::SUCCESS
    }

    unsafe extern "C" fn mem_malloc_callback(
        ptr: *mut *mut c_void,
        requested_len: u32,
        allocated_len: *mut u32,
    ) -> s2n_status_code::Type {
        let layout = if let Some(layout) = layout(requested_len) {
            layout
        } else {
            return s2n_status_code::SUCCESS;
        };
        *ptr = alloc(layout) as *mut _;

        if ptr.is_null() {
            s2n_status_code::FAILURE
        } else {
            *allocated_len = requested_len;
            s2n_status_code::SUCCESS
        }
    }

    unsafe extern "C" fn mem_free_callback(ptr: *mut c_void, len: u32) -> s2n_status_code::Type {
        let layout = if let Some(layout) = layout(len) {
            layout
        } else {
            return s2n_status_code::FAILURE;
        };

        if !ptr.is_null() {
            dealloc(ptr as *mut _, layout);
        }

        s2n_status_code::SUCCESS
    }

    unsafe fn layout(len: u32) -> Option<Layout> {
        // https://linux.die.net/man/3/malloc
        //# The malloc() and calloc() functions return a pointer to the
        //# allocated memory, which is suitably aligned for any built-in
        //# type.
        const ALIGNMENT: usize = size_of::<usize>();

        // * align must not be zero,
        //
        // * align must be a power of two,
        //
        // * size, when rounded up to the nearest multiple of align, must not overflow (i.e., the rounded value must be less than or equal to usize::MAX).

        Layout::from_size_align(len as usize, ALIGNMENT).ok()
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

extern crate alloc;

// Ensure memory is correctly managed in tests
// tests invoked using the checkers::test macro have additional
// memory sanity checks that occur
#[cfg(test)]
#[global_allocator]
static ALLOCATOR: checkers::Allocator = checkers::Allocator::system();

#[macro_use]
pub mod error;

pub mod callbacks;
#[cfg(feature = "unstable-fingerprint")]
pub mod client_hello;
pub mod config;
pub mod connection;
pub mod enums;
pub mod init;
pub mod pool;
pub mod security;

pub use s2n_tls_sys as ffi;

#[cfg(any(feature = "testing", test))]
pub mod testing;


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

//! Utilities to handle reusing connections.
//!
//! Creating a single new connection requires significant
//! memory allocations (about 50-60 KB, according to some tests).
//! Instead of allocating memory for a new connection, existing
//! memory can be reused by calling
//! [Connection::wipe()](`crate::connection::Connection::wipe()).
//!
//! The [`Pool`] trait allows applications to define an
//! [Object pool](https://en.wikipedia.org/wiki/Object_pool_pattern) that
//! wipes and stores connections after they are dropped.
//!
//! We also provide a basic Pool implementation, [`ConfigPool`], that
//! implements the pool as a [VecDeque](`std::collections::VecDeque`)
//! with a fixed maximum size.

use crate::{
    config::Config,
    connection::{Builder, Connection},
    enums::Mode,
    error::Error,
};
use std::{
    collections::VecDeque,
    ops::{Deref, DerefMut},
    sync::{Arc, Mutex},
};

/// A connection produced by a [`Pool`].
///
/// When dropped, returns ownership of the connection to
/// the pool that produced it by calling [`Pool::give`].
#[derive(Debug)]
pub struct PooledConnection<T: Pool = Arc<dyn Pool>> {
    pool: T,
    conn: Option<Connection>,
}

impl<T: Pool> AsRef<Connection> for PooledConnection<T> {
    fn as_ref(&self) -> &Connection {
        self.conn.as_ref().unwrap()
    }
}

impl<T: Pool> AsMut<Connection> for PooledConnection<T> {
    fn as_mut(&mut self) -> &mut Connection {
        self.conn.as_mut().unwrap()
    }
}

impl<T: Pool> Drop for PooledConnection<T> {
    fn drop(&mut self) {
        if let Some(conn) = self.conn.take() {
            self.pool.give(conn);
        }
    }
}

impl<T: Pool> Deref for PooledConnection<T> {
    type Target = Connection;

    fn deref(&self) -> &Self::Target {
        self.conn.as_ref().unwrap()
    }
}

impl<T: Pool> DerefMut for PooledConnection<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.conn.as_mut().unwrap()
    }
}

impl<T: Pool + Clone> PooledConnection<T> {
    pub fn new(pool: &T) -> Result<PooledConnection<T>, Error> {
        pool.take().map(|conn| {
            let conn = Some(conn);
            let pool = pool.clone();
            Self { pool, conn }
        })
    }
}

/// An object pool for wiping and reusing connection memory.
///
/// Minimally, an implementation should call [`Connection::wipe()`]
/// during [`Self::give`].
pub trait Pool {
    fn mode(&self) -> Mode;
    fn take(&self) -> Result<Connection, Error>;
    fn give(&self, conn: Connection);
}

impl Pool for Arc<dyn Pool> {
    fn mode(&self) -> Mode {
        self.as_ref().mode()
    }
    fn take(&self) -> Result<Connection, Error> {
        self.as_ref().take()
    }
    fn give(&self, conn: Connection) {
        self.as_ref().give(conn)
    }
}

impl<T: Pool> Pool for Arc<T> {
    fn mode(&self) -> Mode {
        self.as_ref().mode()
    }
    fn take(&self) -> Result<Connection, Error> {
        self.as_ref().take()
    }
    fn give(&self, conn: Connection) {
        self.as_ref().give(conn)
    }
}

#[derive(Debug)]
pub struct ConfigPool {
    mode: Mode,
    config: Config,
    pool: Mutex<VecDeque<Connection>>,
    max_pool_size: usize,
}

pub type ConfigPoolRef = Arc<ConfigPool>;

/// Builder for [`ConfigPool`].
pub struct ConfigPoolBuilder(ConfigPool);
impl ConfigPoolBuilder {
    pub fn new(mode: Mode, config: Config) -> Self {
        Self(ConfigPool {
            mode,
            config,
            pool: Mutex::new(VecDeque::new()),
            max_pool_size: usize::MAX,
        })
    }

    pub fn set_pool(&mut self, pool: VecDeque<Connection>) -> &mut Self {
        self.0.pool = Mutex::new(pool);
        self
    }

    /// The maximum size of the underlying [`VecDeque`].
    ///
    /// This is NOT the maximum connections that can be created.
    /// When the number of connections created exceeds the `max_pool_size`,
    /// excess reclaimed connections are dropped instead of stored
    /// in the pool.
    pub fn set_max_pool_size(&mut self, max_pool_size: usize) -> &mut Self {
        self.0.max_pool_size = max_pool_size;
        self
    }

    pub fn build(self) -> Arc<ConfigPool> {
        Arc::new(self.0)
    }
}

impl ConfigPool {
    pub fn pool_size(&self) -> usize {
        self.pool.lock().map(|pool| pool.len()).unwrap_or(0)
    }

    pub fn is_poisoned(&self) -> bool {
        self.pool.is_poisoned()
    }
}

impl Pool for ConfigPool {
    fn mode(&self) -> Mode {
        self.mode
    }

    /// Get a connection.
    ///
    /// If connections are available in the pool, one will
    /// be returned. Otherwise, a new connection will be created.
    fn take(&self) -> Result<Connection, Error> {
        let from_pool = match self.pool.lock() {
            Ok(mut pool) => pool.pop_front(),
            Err(_) => None,
        };
        let conn = match from_pool {
            // Wiping a connection doesn't wipe the config,
            // so we don't need to reset the config.
            Some(conn) => conn,
            // Create a new connection with the stored config.
            None => self.config.build_connection(self.mode)?,
        };
        Ok(conn)
    }

    /// Recycle a connection.
    ///
    /// The connection is wiped and returned to the pool
    /// if space is available.
    fn give(&self, mut conn: Connection) {
        let wiped = conn.wipe().is_ok();
        if let Ok(mut pool) = self.pool.lock() {
            if pool.len() < self.max_pool_size && wiped {
                pool.push_back(conn);
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::Config;

    #[test]
    fn config_pool_single_connection() -> Result<(), Box<dyn std::error::Error>> {
        let pool = ConfigPoolBuilder::new(Mode::Server, Config::default()).build();

        // Repeatedly checkout the same connection
        assert_eq!(pool.pool_size(), 0);
        for _ in 0..5 {
            let _conn = PooledConnection::new(&pool)?;
            assert_eq!(pool.pool_size(), 0);
        }
        assert_eq!(pool.pool_size(), 1);
        Ok(())
    }

    #[test]
    fn config_pool_multiple_connections() -> Result<(), Box<dyn std::error::Error>> {
        let pool = ConfigPoolBuilder::new(Mode::Server, Config::default()).build();

        // We need to hold onto connections so that they're not
        // immediately reclaimed by the pool.
        let mut conns = VecDeque::new();

        // Checkout multiple connections
        const COUNT: usize = 25;
        for _ in 0..COUNT {
            conns.push_back(PooledConnection::new(&pool)?);
            assert_eq!(pool.pool_size(), 0);
        }
        assert_eq!(conns.len(), COUNT);

        // Drop all outstanding connections, returning them to the pool
        conns.clear();
        assert_eq!(pool.pool_size(), COUNT);

        // Reuse a subset of the connections
        const SUBSET_COUNT: usize = COUNT / 2;
        for i in 1..=SUBSET_COUNT {
            conns.push_back(PooledConnection::new(&pool)?);
            // The pool should drain as we reuse connections.
            assert_eq!(pool.pool_size(), COUNT - i);
        }
        assert_eq!(conns.len(), SUBSET_COUNT);
        assert_eq!(pool.pool_size(), COUNT - SUBSET_COUNT);

        // Drop all outstanding connections, returning them to the pool
        conns.clear();
        assert_eq!(pool.pool_size(), COUNT);

        Ok(())
    }

    #[test]
    fn config_pool_with_max_size() -> Result<(), Box<dyn std::error::Error>> {
        const POOL_MAX_SIZE: usize = 11;
        let mut pool = ConfigPoolBuilder::new(Mode::Server, Config::default());
        pool.set_max_pool_size(POOL_MAX_SIZE);
        let pool = pool.build();

        // We need to hold onto connections so that they're not
        // immediately reclaimed by the pool.
        let mut conns = VecDeque::new();

        // Create more connections than the pools can hold
        const COUNT: usize = 25;
        for _ in 0..COUNT {
            conns.push_back(PooledConnection::new(&pool)?);
        }
        assert_eq!(conns.len(), COUNT);

        // Drop all outstanding connections, returning them to the pools
        // The pool should now hold its maximum.
        conns.clear();
        assert_eq!(pool.pool_size(), POOL_MAX_SIZE);

        Ok(())
    }

    #[test]
    fn non_generic_pool() -> Result<(), Box<dyn std::error::Error>> {
        let config_pool = ConfigPoolBuilder::new(Mode::Server, Config::default()).build();
        // Note the unweildy type parameters on PooledConnection here.
        let _: PooledConnection<ConfigPoolRef> = PooledConnection::new(&config_pool)?;
        // To avoid specifying the generic type parameters on PooledConnection,
        // the pool can be converted to an Arc<dyn Pool>.
        let pool: Arc<dyn Pool> = config_pool;
        // Note no generic type parameters on PooledConnection here.
        let _: PooledConnection = PooledConnection::new(&pool)?;
        Ok(())
    }

    #[test]
    fn dereferencing_pooled_connection() -> Result<(), Box<dyn std::error::Error>> {
        let config_pool = ConfigPoolBuilder::new(Mode::Server, Config::default()).build();

        let pooled_conn: PooledConnection<ConfigPoolRef> = PooledConnection::new(&config_pool)?;
        let conn = pooled_conn.deref();
        assert_eq!(pooled_conn.config(), conn.config());

        let mut mut_pooled_conn: PooledConnection<ConfigPoolRef> =
            PooledConnection::new(&config_pool)?;
        let waker = futures_test::task::new_count_waker().0;
        mut_pooled_conn.set_waker(Some(&waker))?;
        assert!(mut_pooled_conn.waker().unwrap().will_wake(&waker));

        Ok(())
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::error::Error;
use core::fmt;
use std::ffi::{CStr, CString};

#[derive(Clone, PartialEq)]
enum Context {
    Static(&'static [u8]),
    Owned(CString),
}

#[derive(Clone, PartialEq)]
pub struct Policy(Context);

impl Policy {
    pub(crate) fn as_cstr(&self) -> &CStr {
        match &self.0 {
            Context::Static(x) => unsafe {
                // Safety: Policies are always created with null-terminated strings
                CStr::from_bytes_with_nul_unchecked(x)
            },
            Context::Owned(x) => x.as_c_str(),
        }
    }

    pub fn from_version(version: &str) -> Result<Policy, Error> {
        let cstr = CString::new(version).map_err(|_| Error::INVALID_INPUT)?;
        let context = Context::Owned(cstr);
        Ok(Self(context))
    }
}

impl fmt::Debug for Policy {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_tuple("Policy").field(&self.as_cstr()).finish()
    }
}

macro_rules! policy {
    ($name:ident, $version:expr) => {
        pub const $name: Policy = Policy(Context::Static(concat!($version, "\0").as_bytes()));
    };
}

policy!(DEFAULT, "default");
policy!(DEFAULT_TLS13, "default_tls13");

#[cfg(feature = "pq")]
policy!(TESTING_PQ, "PQ-TLS-1-0-2021-05-26");

pub const ALL_POLICIES: &[Policy] = &[
    DEFAULT,
    DEFAULT_TLS13,
    #[cfg(feature = "pq")]
    TESTING_PQ,
];


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{
    callbacks::VerifyHostNameCallback, config::*, enums::Blinding, security,
    testing::s2n_tls::Harness,
};
use alloc::{collections::VecDeque, sync::Arc};
use bytes::Bytes;
use core::{
    sync::atomic::{AtomicUsize, Ordering},
    task::Poll,
};

pub mod client_hello;
pub mod s2n_tls;

type Error = Box<dyn std::error::Error>;
type Result<T, E = Error> = core::result::Result<T, E>;

pub fn test_error(msg: &str) -> crate::error::Error {
    crate::error::Error::application(msg.into())
}

pub fn assert_test_error(input: Error, msg: &str) {
    let error = input
        .downcast::<crate::error::Error>()
        .expect("Unexpected generic error type");
    if let Some(inner) = error.application_error() {
        assert_eq!(msg, inner.to_string())
    } else {
        panic!("Unexpected known error type");
    }
}

#[derive(Clone)]
pub struct Counter(Arc<AtomicUsize>);
impl Counter {
    fn new() -> Self {
        Counter(Arc::new(AtomicUsize::new(0)))
    }
    pub fn count(&self) -> usize {
        self.0.load(Ordering::Relaxed)
    }
    pub fn increment(&self) {
        self.0.fetch_add(1, Ordering::Relaxed);
    }
}
impl Default for Counter {
    fn default() -> Self {
        Self::new()
    }
}

pub trait Connection: core::fmt::Debug {
    fn poll<Ctx: Context>(&mut self, context: &mut Ctx) -> Poll<Result<()>>;
}

pub trait Context {
    fn receive(&mut self, max_len: Option<usize>) -> Option<Bytes>;
    fn send(&mut self, data: Bytes);
}

#[derive(Debug)]
pub struct Pair<Server: Connection, Client: Connection> {
    pub server: (Server, MemoryContext),
    pub client: (Client, MemoryContext),
    pub max_iterations: usize,
}

impl<Server: Connection, Client: Connection> Pair<Server, Client> {
    /// The number of iterations that will be executed until the handshake exits with an error
    ///
    /// This is to prevent endless looping without making progress on the connection.
    const DEFAULT_ITERATIONS: usize = 100;

    pub fn new(server: Server, client: Client) -> Self {
        Self {
            server: (server, Default::default()),
            client: (client, Default::default()),
            max_iterations: Self::DEFAULT_ITERATIONS,
        }
    }
    pub fn poll(&mut self) -> Poll<Result<()>> {
        assert!(
            self.max_iterations > 0,
            "handshake has iterated too many times: {:#?}",
            self,
        );
        let client_res = self.client.0.poll(&mut self.client.1);
        let server_res = self.server.0.poll(&mut self.server.1);
        self.client.1.transfer(&mut self.server.1);
        self.max_iterations -= 1;
        match (client_res, server_res) {
            (Poll::Ready(client_res), Poll::Ready(server_res)) => {
                client_res?;
                server_res?;
                Ok(()).into()
            }
            (Poll::Ready(client_res), _) => {
                client_res?;
                Poll::Pending
            }
            (_, Poll::Ready(server_res)) => {
                server_res?;
                Poll::Pending
            }
            _ => Poll::Pending,
        }
    }
}

#[derive(Debug, Default)]
pub struct MemoryContext {
    rx: VecDeque<Bytes>,
    tx: VecDeque<Bytes>,
}

impl MemoryContext {
    pub fn transfer(&mut self, other: &mut Self) {
        self.rx.extend(other.tx.drain(..));
        other.rx.extend(self.tx.drain(..));
    }
}

impl Context for MemoryContext {
    fn receive(&mut self, max_len: Option<usize>) -> Option<Bytes> {
        loop {
            let mut chunk = self.rx.pop_front()?;

            if chunk.is_empty() {
                continue;
            }

            let max_len = max_len.unwrap_or(usize::MAX);

            if chunk.len() > max_len {
                self.rx.push_front(chunk.split_off(max_len));
            }

            return Some(chunk);
        }
    }

    fn send(&mut self, data: Bytes) {
        self.tx.push_back(data);
    }
}

pub struct CertKeyPair {
    cert_path: &'static str,
    cert: &'static [u8],
    key: &'static [u8],
}

impl Default for CertKeyPair {
    fn default() -> Self {
        CertKeyPair {
            cert_path: concat!(
                env!("CARGO_MANIFEST_DIR"),
                "/../../../tests/pems/rsa_4096_sha512_client_cert.pem",
            ),
            cert: &include_bytes!("../../../../tests/pems/rsa_4096_sha512_client_cert.pem")[..],
            key: &include_bytes!("../../../../tests/pems/rsa_4096_sha512_client_key.pem")[..],
        }
    }
}

impl CertKeyPair {
    pub fn cert_path(&self) -> &'static str {
        self.cert_path
    }

    pub fn cert(&self) -> &'static [u8] {
        self.cert
    }

    pub fn key(&self) -> &'static [u8] {
        self.key
    }
}

pub struct InsecureAcceptAllCertificatesHandler {}
impl VerifyHostNameCallback for InsecureAcceptAllCertificatesHandler {
    fn verify_host_name(&self, _host_name: &str) -> bool {
        true
    }
}

pub struct RejectAllCertificatesHandler {}
impl VerifyHostNameCallback for RejectAllCertificatesHandler {
    fn verify_host_name(&self, _host_name: &str) -> bool {
        false
    }
}

pub fn build_config(cipher_prefs: &security::Policy) -> Result<crate::config::Config, Error> {
    let builder = config_builder(cipher_prefs)?;
    Ok(builder.build().expect("Unable to build server config"))
}

pub fn config_builder(cipher_prefs: &security::Policy) -> Result<crate::config::Builder, Error> {
    let mut builder = Builder::new();
    let keypair = CertKeyPair::default();
    // Build a config
    builder
        .set_security_policy(cipher_prefs)
        .expect("Unable to set config cipher preferences");
    builder
        .load_pem(keypair.cert(), keypair.key())
        .expect("Unable to load cert/pem");
    builder
        .set_verify_host_callback(InsecureAcceptAllCertificatesHandler {})
        .expect("Unable to set a host verify callback.");
    unsafe {
        builder
            .disable_x509_verification()
            .expect("Unable to disable x509 verification");
    };
    Ok(builder)
}

pub fn tls_pair(config: crate::config::Config) -> Pair<Harness, Harness> {
    // create and configure a server connection
    let mut server = crate::connection::Connection::new_server();
    // some tests check for connection failure so disable blinding to avoid delay
    server.as_mut().set_blinding(Blinding::SelfService).unwrap();
    server
        .set_config(config.clone())
        .expect("Failed to bind config to server connection");
    let server = Harness::new(server);

    // create a client connection
    let mut client = crate::connection::Connection::new_client();
    // some tests check for connection failure so disable blinding to avoid delay
    client.as_mut().set_blinding(Blinding::SelfService).unwrap();
    client
        .set_config(config)
        .expect("Unable to set client config");
    let client = Harness::new(client);

    Pair::new(server, client)
}

pub fn establish_connection(config: crate::config::Config) {
    // create and configure a server connection
    let mut server = crate::connection::Connection::new_server();
    server
        .set_config(config.clone())
        .expect("Failed to bind config to server connection");
    let server = Harness::new(server);

    // create a client connection
    let mut client = crate::connection::Connection::new_client();
    client
        .set_config(config)
        .expect("Unable to set client config");
    let client = Harness::new(client);

    let pair = Pair::new(server, client);
    poll_tls_pair(pair);
}

pub fn poll_tls_pair(mut pair: Pair<Harness, Harness>) -> Pair<Harness, Harness> {
    loop {
        match pair.poll() {
            Poll::Ready(result) => {
                result.unwrap();
                break;
            }
            Poll::Pending => continue,
        }
    }

    pair
}

pub fn poll_tls_pair_result(pair: &mut Pair<Harness, Harness>) -> Result<()> {
    loop {
        match pair.poll() {
            Poll::Ready(result) => return result,
            Poll::Pending => continue,
        }
    }
}