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

//! Support for asynchronous callbacks.
//!
//! The general flow for an async callback is:
//! 1. The application sets FooCallback on the [`crate::config::Config`] with
//!    a method like Config::set_foo_callback.
//! 2. When the underlying C library reaches the trigger for that specific
//!    callback (for example, the ClientHello for [`crate::callbacks::ClientHelloCallback`])
//!    it calls the callback implementation to get a [`ConnectionFuture`].
//! 3. The [`ConnectionFuture`] is stored on the connection. Every time
//!    the handshake is polled, the [`ConnectionFuture`] is polled instead.
//! 4. Once the [`ConnectionFuture`] returns a result, the connection
//!    drops the future and proceeds as usual.

use crate::{connection::Connection, enums::CallbackResult, error::Error};
use core::task::Poll;
use pin_project_lite::pin_project;
use std::pin::Pin;

/// The Future associated with the async connection callback.
///
/// The calling application can provide an instance of [`ConnectionFuture`]
/// when implementing an async callback, eg. [`crate::callbacks::ClientHelloCallback`],
/// if it wants to run an asynchronous operation (disk read, network call).
/// The application can return an error ([Err(Error::application())])
/// to indicate connection failure.
pub trait ConnectionFuture {
    fn poll(
        self: Pin<&mut Self>,
        connection: &mut Connection,
        ctx: &mut core::task::Context,
    ) -> Poll<Result<(), Error>>;
}

pub(crate) type ConnectionFutureResult = Result<Option<Pin<Box<dyn ConnectionFuture>>>, Error>;

// Useful for propagating [`error::Error`] from a C callback back to the Rust application.
pub(crate) struct ErrorFuture {
    error: Option<Error>,
}

impl ConnectionFuture for ErrorFuture {
    fn poll(
        mut self: Pin<&mut Self>,
        _connection: &mut Connection,
        _ctx: &mut core::task::Context,
    ) -> Poll<Result<(), Error>> {
        let err = self.error.take().expect(
            "ErrorFuture should be initialized with Some(error) and a Future should never
            be polled after it returns Poll::Ready",
        );
        Poll::Ready(Err(err))
    }
}

pin_project! {
    /// A wrapper around an optional [`ConnectionFuture`]
    /// which either polls the future or immediately reports success.
    struct OptionalFuture{
        option: Option<Pin<Box<dyn ConnectionFuture>>>,
    }
}

impl OptionalFuture {
    fn new(input: ConnectionFutureResult) -> Self {
        match input {
            Ok(option) => OptionalFuture { option },
            Err(error) => {
                let error = Some(error);
                OptionalFuture {
                    option: Some(Box::pin(ErrorFuture { error })),
                }
            }
        }
    }
}

impl ConnectionFuture for OptionalFuture {
    fn poll(
        mut self: Pin<&mut Self>,
        conn: &mut Connection,
        ctx: &mut core::task::Context,
    ) -> Poll<Result<(), Error>> {
        match self.option.as_mut() {
            Some(future) => future.as_mut().poll(conn, ctx),
            None => Poll::Ready(Ok(())),
        }
    }
}

/// Any work necessary after the callback completes.
//
// We do not expect any callback except [`ClientHelloCallback`] to require MarkDone.
// More recent callbacks follow a different model that doesn't require separate cleanup.
//
// This enum is sufficient while only ClientHello is special-cased, but will not
// scale well. If we need more MarkDone variants, then we should consider a different
// solution, like another stored future.
#[non_exhaustive]
#[derive(PartialEq)]
enum MarkDone {
    ClientHello,
    None,
}

pin_project! {
    // Stores the [`ConnectionFuture`] and associated state.
    pub(crate) struct AsyncCallback {
        #[pin]
        future: OptionalFuture,
        cleanup: MarkDone,
    }
}

impl AsyncCallback {
    pub(crate) fn poll(
        self: Pin<&mut Self>,
        conn: &mut Connection,
        ctx: &mut core::task::Context,
    ) -> Poll<Result<(), Error>> {
        let this = self.project();
        let poll = this.future.poll(conn, ctx);
        if let Poll::Ready(Ok(())) = poll {
            if this.cleanup == &MarkDone::ClientHello {
                conn.mark_client_hello_cb_done()?;
            }
        }
        poll
    }

    pub(crate) fn trigger_client_hello_cb(
        future: ConnectionFutureResult,
        conn: &mut Connection,
    ) -> CallbackResult {
        let future = OptionalFuture::new(future);
        let cleanup = MarkDone::ClientHello;
        let callback = AsyncCallback { future, cleanup };
        conn.set_async_callback(callback);
        CallbackResult::Success
    }

    pub(crate) fn trigger(future: ConnectionFutureResult, conn: &mut Connection) -> CallbackResult {
        let future = OptionalFuture::new(future);
        let cleanup = MarkDone::None;
        let callback = AsyncCallback { future, cleanup };
        conn.set_async_callback(callback);
        CallbackResult::Success
    }
}


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

//! Support for application-implemented ClientHello callbacks.

use crate::{callbacks::*, config::Config, connection::Connection, error::Error};
use core::task::Poll;
use pin_project_lite::pin_project;
use std::{future::Future, pin::Pin};

/// A trait for the callback executed after parsing the ClientHello message.
///
/// Use in conjunction with
/// [config::Builder::set_client_hello_callback](`crate::config::Builder::set_client_hello_callback()`).
pub trait ClientHelloCallback {
    /// The application can return an `Ok(None)` to resolve the callback
    /// synchronously or return an `Ok(Some(ConnectionFuture))` if it wants to
    /// run some asynchronous task before resolving the callback.
    ///
    /// [`ConfigResolver`], which implements [`ConnectionFuture`] can be
    /// returned if the application wants to set a new [`Config`] on the connection.
    ///
    /// If the server_name is used to configure the connection then the application
    /// should call [`Connection::server_name_extension_used()`].
    fn on_client_hello(
        // this method takes an immutable reference to self to prevent the
        // Config from being mutated by one connection and then used in another
        // connection, leading to undefined behavior
        &self,
        connection: &mut Connection,
    ) -> ConnectionFutureResult;
}

// For more information on projection:
// https://doc.rust-lang.org/std/pin/index.html#projections-and-structural-pinning
pin_project! {
    /// An implementation of [`ConnectionFuture`] which resolves the provided
    /// future and sets the config on the [`Connection`].
    pub struct ConfigResolver<F: Future<Output = Result<Config, Error>>> {
        #[pin]
        fut: F,
    }
}

impl<F: Future<Output = Result<Config, Error>>> ConfigResolver<F> {
    pub fn new(fut: F) -> Self {
        ConfigResolver { fut }
    }
}

impl<F: Future<Output = Result<Config, Error>>> ConnectionFuture for ConfigResolver<F> {
    fn poll(
        self: Pin<&mut Self>,
        connection: &mut Connection,
        ctx: &mut core::task::Context,
    ) -> Poll<Result<(), Error>> {
        let this = self.project();
        let config = match this.fut.poll(ctx) {
            Poll::Ready(config) => config?,
            Poll::Pending => return Poll::Pending,
        };

        connection.set_config(config)?;

        Poll::Ready(Ok(()))
    }
}


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

use crate::{
    callbacks::*,
    connection::Connection,
    enums::{HashAlgorithm, Mode, SignatureAlgorithm},
    error::{Error, Fallible},
    ffi::*,
};
use std::{pin::Pin, ptr::NonNull};

#[non_exhaustive]
#[derive(Debug)]
pub enum OperationType {
    Decrypt,
    Sign(SignatureAlgorithm, HashAlgorithm),
}

pub struct PrivateKeyOperation {
    raw: NonNull<s2n_async_pkey_op>,
    kind: OperationType,
}

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

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

impl PrivateKeyOperation {
    pub(crate) fn try_from_cb(
        conn: &Connection,
        op_ptr: *mut s2n_async_pkey_op,
    ) -> Result<Self, Error> {
        let mut raw_kind = 0;
        unsafe { s2n_async_pkey_op_get_op_type(op_ptr, &mut raw_kind) }.into_result()?;

        let kind = match raw_kind {
            s2n_async_pkey_op_type::SIGN => {
                let sig_alg = match conn.mode() {
                    Mode::Client => conn
                        .selected_client_signature_algorithm()?
                        .ok_or(Error::INVALID_INPUT)?,
                    Mode::Server => conn.selected_signature_algorithm()?,
                };
                let hash_alg = match conn.mode() {
                    Mode::Client => conn
                        .selected_client_hash_algorithm()?
                        .ok_or(Error::INVALID_INPUT)?,
                    Mode::Server => conn.selected_hash_algorithm()?,
                };
                OperationType::Sign(sig_alg, hash_alg)
            }
            s2n_async_pkey_op_type::DECRYPT => OperationType::Decrypt,
            _ => return Err(Error::INVALID_INPUT),
        };

        let raw = NonNull::new(op_ptr).ok_or(Error::INVALID_INPUT)?;
        Ok(PrivateKeyOperation { raw, kind })
    }

    /// Do we need to sign or decrypt with the private key?
    pub fn kind(&self) -> Result<&OperationType, Error> {
        Ok(&self.kind)
    }

    /// The size of the slice returned by [`input()`]
    pub fn input_size(&self) -> Result<usize, Error> {
        let mut size = 0;
        unsafe { s2n_async_pkey_op_get_input_size(self.raw.as_ptr(), &mut size) }.into_result()?;
        size.try_into().map_err(|_| Error::INVALID_INPUT)
    }

    /// Provides the input for the operation.
    ///
    /// If this is an [`OperationType::Sign`] operation, then this input has
    /// already been hashed and is the resultant digest.
    pub fn input(&self, buf: &mut [u8]) -> Result<(), Error> {
        let buf_len: u32 = buf.len().try_into().map_err(|_| Error::INVALID_INPUT)?;
        let buf_ptr = buf.as_ptr() as *mut u8;
        unsafe { s2n_async_pkey_op_get_input(self.raw.as_ptr(), buf_ptr, buf_len) }
            .into_result()?;
        Ok(())
    }

    /// Sets the output of the operation
    pub fn set_output(self, conn: &mut Connection, buf: &[u8]) -> Result<(), Error> {
        let buf_len: u32 = buf.len().try_into().map_err(|_| Error::INVALID_INPUT)?;
        let buf_ptr = buf.as_ptr() as *const u8;
        unsafe {
            s2n_async_pkey_op_set_output(self.raw.as_ptr(), buf_ptr, buf_len).into_result()?;
            s2n_async_pkey_op_apply(self.raw.as_ptr(), conn.as_ptr()).into_result()?;
        }
        Ok(())
    }
}

impl Drop for PrivateKeyOperation {
    fn drop(&mut self) {
        unsafe {
            let _ = s2n_async_pkey_op_free(self.raw.as_ptr());
        }
    }
}

pub trait PrivateKeyCallback {
    fn handle_operation(
        &self,
        connection: &mut Connection,
        operation: PrivateKeyOperation,
    ) -> Result<Option<Pin<Box<dyn ConnectionFuture>>>, Error>;
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        config, connection, error, security, testing,
        testing::{s2n_tls::*, *},
    };
    use core::task::{Poll, Waker};
    use futures_test::task::new_count_waker;
    use openssl::{ec::EcKey, ecdsa::EcdsaSig};

    type Error = Box<dyn std::error::Error>;

    const KEY: &[u8] = include_bytes!(concat!(
        env!("CARGO_MANIFEST_DIR"),
        "/../../../tests/pems/ecdsa_p384_pkcs1_key.pem"
    ));
    const CERT: &[u8] = include_bytes!(concat!(
        env!("CARGO_MANIFEST_DIR"),
        "/../../../tests/pems/ecdsa_p384_pkcs1_cert.pem"
    ));

    fn new_pair<T>(
        callback: T,
        waker: Waker,
    ) -> Result<Pair<s2n_tls::Harness, s2n_tls::Harness>, Error>
    where
        T: 'static + PrivateKeyCallback,
    {
        let config = {
            let mut config = config::Builder::new();
            config.set_security_policy(&security::DEFAULT_TLS13)?;
            config.load_public_pem(CERT)?;
            config.set_private_key_callback(callback)?;
            // Our test certificates are untrusted, but disabling certificate
            // verification does not affect handshake signatures.
            unsafe { config.disable_x509_verification() }?;
            config.build()?
        };

        let server = {
            let mut server = connection::Connection::new_server();
            server.set_config(config.clone())?;
            server.set_waker(Some(&waker))?;
            Harness::new(server)
        };

        let client = {
            let mut client = connection::Connection::new_client();
            client.set_config(config)?;
            Harness::new(client)
        };

        Ok(Pair::new(server, client))
    }

    fn ecdsa_sign(
        op: PrivateKeyOperation,
        conn: &mut connection::Connection,
        key: &[u8],
    ) -> Result<(), error::Error> {
        match op.kind()? {
            OperationType::Sign(SignatureAlgorithm::ECDSA, _) => {
                let in_buf_size = op.input_size()?;
                let mut in_buf = vec![0; in_buf_size];
                op.input(&mut in_buf)?;

                let key =
                    EcKey::private_key_from_pem(key).expect("Failed to create EcKey from pem");
                let sig = EcdsaSig::sign(&in_buf, &key).expect("Failed to sign input");
                let out = sig.to_der().expect("Failed to convert signature to der");

                op.set_output(conn, &out)?;
            }
            _ => panic!("Unexpected pkey operation"),
        }
        Ok(())
    }

    #[test]
    fn sync_offload_success() -> Result<(), Error> {
        struct TestPkeyCallback(Counter);
        impl PrivateKeyCallback for TestPkeyCallback {
            fn handle_operation(
                &self,
                conn: &mut connection::Connection,
                op: PrivateKeyOperation,
            ) -> Result<Option<Pin<Box<dyn ConnectionFuture>>>, error::Error> {
                self.0.increment();
                ecdsa_sign(op, conn, KEY)?;
                Ok(None)
            }
        }

        let (waker, wake_count) = new_count_waker();
        let counter = testing::Counter::default();
        let callback = TestPkeyCallback(counter.clone());
        let pair = new_pair(callback, waker)?;

        assert_eq!(counter.count(), 0);
        assert_eq!(wake_count, 0);
        poll_tls_pair(pair);
        assert_eq!(counter.count(), 1);
        assert_eq!(wake_count, 0);

        Ok(())
    }

    #[test]
    fn async_offload_success() -> Result<(), Error> {
        const POLL_COUNT: usize = 10;

        struct TestPkeyFuture {
            counter: usize,
            op: Option<PrivateKeyOperation>,
        }
        impl ConnectionFuture for TestPkeyFuture {
            fn poll(
                mut self: Pin<&mut Self>,
                conn: &mut connection::Connection,
                ctx: &mut core::task::Context,
            ) -> Poll<Result<(), error::Error>> {
                ctx.waker().wake_by_ref();
                self.counter += 1;
                if self.counter < POLL_COUNT {
                    Poll::Pending
                } else if let Some(op) = self.op.take() {
                    Poll::Ready(ecdsa_sign(op, conn, KEY))
                } else {
                    Poll::Ready(Err(error::Error::application(
                        "missing pkey operation".into(),
                    )))
                }
            }
        }

        struct TestPkeyCallback(Counter);
        impl PrivateKeyCallback for TestPkeyCallback {
            fn handle_operation(
                &self,
                _conn: &mut connection::Connection,
                op: PrivateKeyOperation,
            ) -> Result<Option<Pin<Box<dyn ConnectionFuture>>>, error::Error> {
                self.0.increment();
                let future = TestPkeyFuture {
                    counter: 0,
                    op: Some(op),
                };
                Ok(Some(Box::pin(future)))
            }
        }

        let (waker, wake_count) = new_count_waker();
        let counter = testing::Counter::default();
        let callback = TestPkeyCallback(counter.clone());
        let pair = new_pair(callback, waker)?;

        assert_eq!(counter.count(), 0);
        assert_eq!(wake_count, 0);
        poll_tls_pair(pair);
        assert_eq!(counter.count(), 1);
        assert_eq!(wake_count, POLL_COUNT);

        Ok(())
    }

    #[test]
    fn sync_failure() -> Result<(), Error> {
        const ERROR: &str = "sync_failure error";

        struct TestPkeyCallback(Counter);
        impl PrivateKeyCallback for TestPkeyCallback {
            fn handle_operation(
                &self,
                _conn: &mut connection::Connection,
                _op: PrivateKeyOperation,
            ) -> Result<Option<Pin<Box<dyn ConnectionFuture>>>, error::Error> {
                self.0.increment();
                Err(testing::test_error(ERROR))
            }
        }

        let (waker, wake_count) = new_count_waker();
        let counter = testing::Counter::default();
        let callback = TestPkeyCallback(counter.clone());
        let mut pair = new_pair(callback, waker)?;

        assert_eq!(counter.count(), 0);
        assert_eq!(wake_count, 0);
        let result = poll_tls_pair_result(&mut pair);
        assert_eq!(counter.count(), 1);
        assert_eq!(wake_count, 0);

        match result {
            Ok(_) => panic!("Handshake unexpectedly succeeded"),
            Err(e) => testing::assert_test_error(e, ERROR),
        };
        Ok(())
    }

    #[test]
    fn async_failure() -> Result<(), Error> {
        const POLL_COUNT: usize = 10;
        const ERROR: &str = "async_failure error";

        struct TestPkeyFuture {
            counter: usize,
            _op: PrivateKeyOperation,
        }
        impl ConnectionFuture for TestPkeyFuture {
            fn poll(
                mut self: Pin<&mut Self>,
                _conn: &mut connection::Connection,
                ctx: &mut core::task::Context,
            ) -> Poll<Result<(), error::Error>> {
                ctx.waker().wake_by_ref();
                self.counter += 1;
                if self.counter < POLL_COUNT {
                    Poll::Pending
                } else {
                    Poll::Ready(Err(testing::test_error(ERROR)))
                }
            }
        }

        struct TestPkeyCallback(Counter);
        impl PrivateKeyCallback for TestPkeyCallback {
            fn handle_operation(
                &self,
                _conn: &mut connection::Connection,
                _op: PrivateKeyOperation,
            ) -> Result<Option<Pin<Box<dyn ConnectionFuture>>>, error::Error> {
                self.0.increment();
                let future = TestPkeyFuture { counter: 0, _op };
                Ok(Some(Box::pin(future)))
            }
        }

        let (waker, wake_count) = new_count_waker();
        let counter = testing::Counter::default();
        let callback = TestPkeyCallback(counter.clone());
        let mut pair = new_pair(callback, waker)?;

        assert_eq!(counter.count(), 0);
        assert_eq!(wake_count, 0);
        let result = poll_tls_pair_result(&mut pair);
        assert_eq!(counter.count(), 1);
        assert_eq!(wake_count, POLL_COUNT);

        match result {
            Ok(_) => panic!("Handshake unexpectedly succeeded"),
            Err(e) => testing::assert_test_error(e, ERROR),
        };
        Ok(())
    }
}


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

use crate::{
    config::Config,
    connection::Connection,
    enums::Mode,
    error::Error,
    pool::{Pool, PooledConnection},
};

/// A trait indicating that a structure can produce connections.
pub trait Builder: Clone {
    type Output: AsMut<Connection> + AsRef<Connection>;
    fn build_connection(&self, mode: Mode) -> Result<Self::Output, Error>;
}

/// Produces new connections with the given Config set.
impl Builder for Config {
    type Output = Connection;
    fn build_connection(&self, mode: Mode) -> Result<Self::Output, Error> {
        let mut conn = Connection::new(mode);
        conn.set_config(self.clone())?;
        Ok(conn)
    }
}

/// Produces new connections from a pool of reuseable connections.
impl<T: Pool + Clone> Builder for T {
    type Output = PooledConnection<T>;
    fn build_connection(&self, mode: Mode) -> Result<Self::Output, Error> {
        if mode == self.mode() {
            Ok(PooledConnection::new(self)?)
        } else {
            Err(Error::INVALID_INPUT)
        }
    }
}

/// Produces new connections from a builder, then modifies them.
///
/// Can be used to apply connection-level config, for example
/// when using a [`crate::pool::ConfigPool`].
#[derive(Clone)]
pub struct ModifiedBuilder<F, B: Builder>
where
    F: Fn(&mut Connection) -> Result<&mut Connection, Error> + Clone,
{
    builder: B,
    modifier: F,
}

impl<F, B: Builder> ModifiedBuilder<F, B>
where
    F: Fn(&mut Connection) -> Result<&mut Connection, Error> + Clone,
{
    pub fn new(builder: B, modifier: F) -> Self {
        Self { builder, modifier }
    }
}

impl<F, B: Builder> Builder for ModifiedBuilder<F, B>
where
    F: Fn(&mut Connection) -> Result<&mut Connection, Error> + Clone,
{
    type Output = B::Output;
    fn build_connection(&self, mode: Mode) -> Result<Self::Output, Error> {
        let mut conn = self.builder.build_connection(mode)?;
        (self.modifier)(conn.as_mut())?;
        Ok(conn)
    }
}

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

    #[test]
    fn config_builder() -> Result<(), Box<dyn std::error::Error>> {
        let config = Config::default();
        let conn = config.build_connection(Mode::Server)?;
        assert_eq!(conn.config(), Some(config));
        Ok(())
    }

    #[test]
    fn pool_builder() -> Result<(), Box<dyn std::error::Error>> {
        let config = Config::default();
        let pool = ConfigPoolBuilder::new(Mode::Server, config.clone()).build();
        let conn = pool.build_connection(Mode::Server)?;
        assert_eq!(conn.as_ref().config(), Some(config));
        Ok(())
    }

    #[test]
    fn modified_builder() -> Result<(), Box<dyn std::error::Error>> {
        let config_a = Config::default();
        let config_b = Config::default();
        assert!(config_a != config_b);

        let builder =
            ModifiedBuilder::new(config_a.clone(), |conn| conn.set_config(config_b.clone()));

        let conn = builder.build_connection(Mode::Server)?;
        assert!(conn.config() != Some(config_a));
        assert_eq!(conn.config(), Some(config_b));
        Ok(())
    }
}


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

use crate::{
    callbacks::{ClientHelloCallback, ConnectionFuture},
    error, security,
};
use alloc::sync::Arc;
use core::{sync::atomic::Ordering, task::Poll};
use s2n_tls_sys::{s2n_client_hello_has_extension, s2n_connection_get_client_hello};
use std::{fmt, io, pin::Pin, sync::atomic::AtomicUsize};

// The Future returned by MockClientHelloHandler.
//
// An instance of this Future is stored on the connection and
// polled to make progress in the async client_hello_callback
pub struct MockClientHelloFuture {
    require_pending_count: usize,
    invoked: Arc<AtomicUsize>,
}

impl ConnectionFuture for MockClientHelloFuture {
    fn poll(
        self: Pin<&mut Self>,
        connection: &mut crate::connection::Connection,
        _ctx: &mut core::task::Context,
    ) -> Poll<Result<(), error::Error>> {
        if self.invoked.fetch_add(1, Ordering::SeqCst) < self.require_pending_count {
            // confirm the callback can access the waker
            connection.waker().unwrap().wake_by_ref();
            return Poll::Pending;
        }

        // Test that the config can be changed
        connection
            .set_config(super::build_config(&security::DEFAULT_TLS13).unwrap())
            .unwrap();

        // Test that server_name_extension_used can be invoked
        connection.server_name_extension_used();

        Poll::Ready(Ok(()))
    }
}

#[derive(Clone)]
pub struct MockClientHelloHandler {
    require_pending_count: usize,
    pub invoked: Arc<AtomicUsize>,
}

impl MockClientHelloHandler {
    pub fn new(require_pending_count: usize) -> Self {
        Self {
            require_pending_count,
            invoked: Arc::new(AtomicUsize::new(0)),
        }
    }
}

impl ClientHelloCallback for MockClientHelloHandler {
    fn on_client_hello(
        &self,
        _connection: &mut crate::connection::Connection,
    ) -> Result<Option<Pin<Box<dyn ConnectionFuture>>>, crate::error::Error> {
        let fut = MockClientHelloFuture {
            require_pending_count: self.require_pending_count,
            invoked: self.invoked.clone(),
        };

        // returning `Some` indicates that the client_hello callback is
        // not yet finished and that the supplied MockClientHelloFuture
        // needs to be `poll`ed to make progress.
        Ok(Some(Box::pin(fut)))
    }
}

// A ClientHelloCallback which returns a Asynchronous task, which
// eventually returns an error.
#[derive(Default)]
pub struct FailingCHHandler;

impl ClientHelloCallback for FailingCHHandler {
    fn on_client_hello(
        &self,
        _connection: &mut crate::connection::Connection,
    ) -> Result<Option<Pin<Box<dyn ConnectionFuture>>>, error::Error> {
        let io_error = io::Error::new(io::ErrorKind::Other, CustomError);
        Err(crate::error::Error::application(Box::new(io_error)))
    }
}

#[derive(Default)]
pub struct FailingAsyncCHHandler;
impl ClientHelloCallback for FailingAsyncCHHandler {
    fn on_client_hello(
        &self,
        _connection: &mut crate::connection::Connection,
    ) -> Result<Option<Pin<Box<dyn ConnectionFuture>>>, error::Error> {
        let fut = FailingCHFuture::default();
        Ok(Some(Box::pin(fut)))
    }
}

// A ClientHelloCallback which returns a synchronous error.
#[derive(Default)]
struct FailingCHFuture {
    pub invoked: Arc<AtomicUsize>,
}

impl ConnectionFuture for FailingCHFuture {
    fn poll(
        self: Pin<&mut Self>,
        connection: &mut crate::connection::Connection,
        _ctx: &mut core::task::Context,
    ) -> Poll<Result<(), error::Error>> {
        if self.invoked.fetch_add(1, Ordering::SeqCst) < 1 {
            // confirm the callback can access the waker
            connection.waker().unwrap().wake_by_ref();
            return Poll::Pending;
        }

        let io_error = io::Error::new(io::ErrorKind::Other, CustomError);
        let ret = Err(crate::error::Error::application(Box::new(io_error)));
        Poll::Ready(ret)
    }
}

impl Drop for FailingCHFuture {
    // return pending once to simulate the async nature of the future and
    // improve test coverage
    fn drop(&mut self) {
        assert!(self.invoked.load(Ordering::SeqCst) >= 1);
    }
}
/// A client hello handler that asserts that the extension with the given
/// IANA code is either present or not present in the client hello
pub struct HasExtensionClientHelloHandler {
    pub extension_iana: u16,
    pub extension_expected: bool,
}

impl ClientHelloCallback for HasExtensionClientHelloHandler {
    fn on_client_hello(
        &self,
        connection: &mut crate::connection::Connection,
    ) -> Result<Option<Pin<Box<dyn ConnectionFuture>>>, error::Error> {
        let mut exists = false;

        unsafe {
            let client_hello = s2n_connection_get_client_hello(connection.as_ptr());
            s2n_client_hello_has_extension(client_hello, self.extension_iana, &mut exists as _);
        }

        if self.extension_expected {
            assert!(
                exists,
                "Extension {} was not found in the client hello",
                self.extension_iana
            );
        } else {
            assert!(
                !exists,
                "Unexpected extension {} found in the client hello",
                self.extension_iana
            )
        }

        Ok(None)
    }
}

#[derive(Debug)]
pub struct CustomError;

impl std::error::Error for CustomError {}
impl fmt::Display for CustomError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "custom error")
    }
}


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

use crate::{
    connection::Connection,
    testing::{Context, Error, Result},
};
use bytes::BytesMut;
use core::task::Poll;
use libc::c_void;
use s2n_tls_sys::s2n_status_code::Type as s2n_status_code;

const SEND_BUFFER_CAPACITY: usize = 4096;

#[derive(Debug)]
pub struct Harness {
    connection: Connection,
    send_buffer: BytesMut,
    handshake_done: bool,
    // TODO add a size
}

impl Harness {
    pub fn new(connection: Connection) -> Self {
        Self {
            connection,
            send_buffer: BytesMut::new(),
            handshake_done: false,
        }
    }

    pub fn connection(&self) -> &Connection {
        &self.connection
    }

    pub fn connection_mut(&mut self) -> &mut Connection {
        &mut self.connection
    }
}

impl super::Connection for Harness {
    fn poll<Ctx: Context>(&mut self, context: &mut Ctx) -> Poll<Result<()>> {
        let mut callback: Callback<Ctx> = Callback {
            context,
            err: None,
            send_buffer: &mut self.send_buffer,
        };

        unsafe {
            // Safety: the callback struct must live as long as the callbacks are
            // set on on the connection
            callback.set(&mut self.connection);
        }

        let result = self.connection.poll_negotiate().map_ok(|_| ());

        callback.unset(&mut self.connection)?;

        match result {
            Poll::Ready(Ok(_)) => {
                if !self.handshake_done {
                    self.handshake_done = true;
                }
                Ok(()).into()
            }
            Poll::Ready(Err(err)) => Err(err.into()).into(),
            Poll::Pending => Poll::Pending,
        }
    }
}

struct Callback<'a, T> {
    pub context: &'a mut T,
    pub err: Option<Error>,
    pub send_buffer: &'a mut BytesMut,
}

impl<'a, T: 'a + Context> Callback<'a, T> {
    unsafe fn set(&mut self, connection: &mut Connection) {
        let context = self as *mut Self as *mut c_void;

        // We use unwrap here since s2n-tls will just check if connection is not null
        connection.set_send_callback(Some(Self::send_cb)).unwrap();
        connection.set_send_context(context).unwrap();
        connection
            .set_receive_callback(Some(Self::recv_cb))
            .unwrap();
        connection.set_receive_context(context).unwrap();
    }

    /// Removes all of the callback and context pointers from the connection
    pub fn unset(mut self, connection: &mut Connection) -> Result<()> {
        unsafe {
            unsafe extern "C" fn send_cb(
                _context: *mut c_void,
                _data: *const u8,
                _len: u32,
            ) -> s2n_status_code {
                -1
            }

            unsafe extern "C" fn recv_cb(
                _context: *mut c_void,
                _data: *mut u8,
                _len: u32,
            ) -> s2n_status_code {
                -1
            }

            // We use unwrap here since s2n-tls will just check if connection is not null
            connection.set_send_callback(Some(send_cb)).unwrap();
            connection.set_send_context(core::ptr::null_mut()).unwrap();
            connection.set_receive_callback(Some(recv_cb)).unwrap();
            connection
                .set_receive_context(core::ptr::null_mut())
                .unwrap();

            // Flush the send buffer before returning to the connection
            self.flush();

            if let Some(err) = self.err {
                return Err(err);
            }

            Ok(())
        }
    }

    unsafe extern "C" fn send_cb(
        context: *mut c_void,
        data: *const u8,
        len: u32,
    ) -> s2n_status_code {
        let context = &mut *(context as *mut Self);
        let data = core::slice::from_raw_parts(data, len as _);
        context.on_write(data) as _
    }

    /// Called when sending data
    fn on_write(&mut self, data: &[u8]) -> usize {
        // If this write would cause the current send buffer to reallocate,
        // we should flush and create a new send buffer.
        let remaining_capacity = self.send_buffer.capacity() - self.send_buffer.len();

        if remaining_capacity < data.len() {
            // Flush the send buffer before reallocating it
            self.flush();

            // ensure we only do one allocation for this write
            let len = SEND_BUFFER_CAPACITY.max(data.len());

            debug_assert!(
                self.send_buffer.is_empty(),
                "dropping a send buffer with data will result in data loss"
            );
            *self.send_buffer = BytesMut::with_capacity(len);
        }

        // Write the current data to the send buffer
        //
        // NOTE: we don't immediately flush to the context since s2n-tls may do
        //       several small writes in a row.
        self.send_buffer.extend_from_slice(data);

        data.len()
    }

    /// Flushes the send buffer into the context
    fn flush(&mut self) {
        if !self.send_buffer.is_empty() {
            let chunk = self.send_buffer.split().freeze();
            self.context.send(chunk);
        }
    }

    /// The function s2n-tls calls when it wants to receive data
    unsafe extern "C" fn recv_cb(context: *mut c_void, data: *mut u8, len: u32) -> s2n_status_code {
        let context = &mut *(context as *mut Self);
        let data = core::slice::from_raw_parts_mut(data, len as _);
        match context.on_read(data) {
            0 => {
                // https://aws.github.io/s2n-tls/doxygen/s2n_8h.html#a699fd9e05a8e8163811db6cab01af973
                // s2n-tls wants us to set the global errno to signal blocked
                errno::set_errno(errno::Errno(libc::EWOULDBLOCK));
                -1
            }
            len => len as _,
        }
    }

    /// Called when receiving data
    fn on_read(&mut self, data: &mut [u8]) -> usize {
        let max_len = Some(data.len());

        // TODO: loop until data buffer is full.
        if let Some(chunk) = self.context.receive(max_len) {
            let len = chunk.len();
            data[..len].copy_from_slice(&chunk);
            len
        } else {
            0
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::{
        callbacks::{ClientHelloCallback, ConnectionFuture},
        enums::ClientAuthType,
        testing::{client_hello::*, s2n_tls::*, *},
    };
    use alloc::sync::Arc;
    use core::sync::atomic::Ordering;
    use futures_test::task::{new_count_waker, noop_waker};
    use std::{fs, path::Path, pin::Pin, sync::atomic::AtomicUsize};

    #[test]
    fn handshake_default() {
        let config = build_config(&security::DEFAULT).unwrap();
        establish_connection(config);
    }

    #[test]
    fn handshake_default_tls13() {
        let config = build_config(&security::DEFAULT_TLS13).unwrap();
        establish_connection(config)
    }

    #[test]
    fn default_config_and_clone_interaction() -> Result<(), Error> {
        let config = build_config(&security::DEFAULT_TLS13)?;
        assert_eq!(config.test_get_refcount()?, 1);
        {
            // Create new connection.
            let mut server = crate::connection::Connection::new_server();
            // Can't retrieve default config.
            assert!(server.config().is_none());
            // Custom config reference count doesn't change.
            assert_eq!(config.test_get_refcount()?, 1);

            // Set custom config on connection.
            server.set_config(config.clone())?;
            // Can retrieve custom config.
            assert!(server.config().is_some());
            // Custom config now referenced once more.
            assert_eq!(config.test_get_refcount()?, 2);

            // Create new connection.
            let mut client = crate::connection::Connection::new_client();
            // Can't retrieve default config.
            assert!(client.config().is_none());
            // Custom config reference count doesn't change.
            assert_eq!(config.test_get_refcount()?, 2);

            // Set custom config on connection.
            client.set_config(config.clone())?;
            // Can retrieve custom config.
            assert!(client.config().is_some());
            // Custom config now referenced once more.
            assert_eq!(config.test_get_refcount()?, 3);

            // drop all the clones
        }
        assert_eq!(config.test_get_refcount()?, 1);
        Ok(())
    }

    #[test]
    fn set_config_multiple_times() -> Result<(), Error> {
        let config = build_config(&security::DEFAULT_TLS13)?;
        assert_eq!(config.test_get_refcount()?, 1);

        let mut server = crate::connection::Connection::new_server();
        assert_eq!(config.test_get_refcount()?, 1);

        // call set_config once
        server.set_config(config.clone())?;
        assert_eq!(config.test_get_refcount()?, 2);
        assert!(server.config().is_some());

        // calling set_config multiple times works since we drop the previous config
        server.set_config(config.clone())?;
        assert_eq!(config.test_get_refcount()?, 2);
        assert!(server.config().is_some());
        Ok(())
    }

    #[test]
    fn connnection_waker() {
        let config = build_config(&security::DEFAULT_TLS13).unwrap();
        assert_eq!(config.test_get_refcount().unwrap(), 1);

        let mut server = crate::connection::Connection::new_server();
        server.set_config(config).unwrap();

        assert!(server.waker().is_none());

        let (waker, wake_count) = new_count_waker();
        server.set_waker(Some(&waker)).unwrap();
        assert!(server.waker().is_some());

        server.set_waker(None).unwrap();
        assert!(server.waker().is_none());

        assert_eq!(wake_count, 0);
    }

    #[test]
    fn failing_client_hello_callback_sync() -> Result<(), Error> {
        let (waker, wake_count) = new_count_waker();
        let handle = FailingCHHandler::default();
        let config = {
            let mut config = config_builder(&security::DEFAULT_TLS13)?;
            config.set_client_hello_callback(handle)?;
            config.build()?
        };

        let server = {
            // create and configure a server connection
            let mut server = crate::connection::Connection::new_server();
            server.set_config(config.clone())?;
            server.set_waker(Some(&waker))?;
            Harness::new(server)
        };

        let client = {
            // create a client connection
            let mut client = crate::connection::Connection::new_client();
            client.set_config(config)?;
            Harness::new(client)
        };

        let mut pair = Pair::new(server, client);
        loop {
            match pair.poll() {
                Poll::Ready(result) => {
                    let err = result.expect_err("handshake should fail");

                    // the underlying error should be the custom error the application provided
                    let s2n_err = err.downcast_ref::<crate::error::Error>().unwrap();
                    let app_err = s2n_err.application_error().unwrap();
                    let io_err = app_err.downcast_ref::<std::io::Error>().unwrap();
                    let _custom_err = io_err
                        .get_ref()
                        .unwrap()
                        .downcast_ref::<CustomError>()
                        .unwrap();
                    break;
                }
                Poll::Pending => continue,
            }
        }
        assert_eq!(wake_count, 0);

        Ok(())
    }

    #[test]
    fn failing_client_hello_callback_async() -> Result<(), Error> {
        let (waker, wake_count) = new_count_waker();
        let handle = FailingAsyncCHHandler::default();
        let config = {
            let mut config = config_builder(&security::DEFAULT_TLS13)?;
            config.set_client_hello_callback(handle)?;
            config.build()?
        };

        let server = {
            // create and configure a server connection
            let mut server = crate::connection::Connection::new_server();
            server.set_config(config.clone())?;
            server.set_waker(Some(&waker))?;
            Harness::new(server)
        };

        let client = {
            // create a client connection
            let mut client = crate::connection::Connection::new_client();
            client.set_config(config)?;
            Harness::new(client)
        };

        let mut pair = Pair::new(server, client);
        loop {
            match pair.poll() {
                Poll::Ready(result) => {
                    let err = result.expect_err("handshake should fail");

                    // the underlying error should be the custom error the application provided
                    let s2n_err = err.downcast_ref::<crate::error::Error>().unwrap();
                    let app_err = s2n_err.application_error().unwrap();
                    let io_err = app_err.downcast_ref::<std::io::Error>().unwrap();
                    let _custom_err = io_err
                        .get_ref()
                        .unwrap()
                        .downcast_ref::<CustomError>()
                        .unwrap();
                    break;
                }
                Poll::Pending => continue,
            }
        }
        // assert that the future is async returned Poll::Pending once
        assert_eq!(wake_count, 1);

        Ok(())
    }

    #[test]
    fn client_hello_callback_async() -> Result<(), Error> {
        let (waker, wake_count) = new_count_waker();
        let require_pending_count = 10;
        let handle = MockClientHelloHandler::new(require_pending_count);
        let config = {
            let mut config = config_builder(&security::DEFAULT_TLS13)?;
            config.set_client_hello_callback(handle.clone())?;
            // multiple calls to set_client_hello_callback should succeed
            config.set_client_hello_callback(handle.clone())?;
            config.build()?
        };

        let server = {
            // create and configure a server connection
            let mut server = crate::connection::Connection::new_server();
            server.set_config(config.clone())?;
            server.set_waker(Some(&waker))?;
            Harness::new(server)
        };

        let client = {
            // create a client connection
            let mut client = crate::connection::Connection::new_client();
            client.set_config(config)?;
            Harness::new(client)
        };

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

        poll_tls_pair(pair);
        // confirm that the callback returned Pending `require_pending_count` times
        assert_eq!(wake_count, require_pending_count);
        // confirm that the final invoked count is +1 more than `require_pending_count`
        assert_eq!(
            handle.invoked.load(Ordering::SeqCst),
            require_pending_count + 1
        );

        Ok(())
    }
    #[test]
    fn client_hello_callback_sync() -> Result<(), Error> {
        let (waker, wake_count) = new_count_waker();
        #[derive(Clone)]
        struct ClientHelloSyncCallback(Arc<AtomicUsize>);
        impl ClientHelloSyncCallback {
            fn new() -> Self {
                ClientHelloSyncCallback(Arc::new(AtomicUsize::new(0)))
            }
            fn count(&self) -> usize {
                self.0.load(Ordering::Relaxed)
            }
        }
        impl ClientHelloCallback for ClientHelloSyncCallback {
            fn on_client_hello(
                &self,
                connection: &mut crate::connection::Connection,
            ) -> Result<Option<Pin<Box<dyn ConnectionFuture>>>, crate::error::Error> {
                // Test that the config can be changed
                connection
                    .set_config(build_config(&security::DEFAULT_TLS13).unwrap())
                    .unwrap();

                // Test that server_name_extension_used can be invoked
                connection.server_name_extension_used();

                self.0.fetch_add(1, Ordering::Relaxed);

                // returning `None` indicates that the client_hello callback is
                // finished and the handshake can proceed.
                Ok(None)
            }
        }
        let callback = ClientHelloSyncCallback::new();

        let config = {
            let mut config = config_builder(&security::DEFAULT_TLS13)?;
            config.set_client_hello_callback(callback.clone())?;
            config.build()?
        };

        let server = {
            // create and configure a server connection
            let mut server = crate::connection::Connection::new_server();
            server.set_config(config.clone())?;
            server.set_waker(Some(&waker))?;
            Harness::new(server)
        };

        let client = {
            // create a client connection
            let mut client = crate::connection::Connection::new_client();
            client.set_config(config)?;
            Harness::new(client)
        };

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

        assert_eq!(callback.count(), 0);
        poll_tls_pair(pair);
        assert_eq!(callback.count(), 1);
        assert_eq!(wake_count, 0);
        Ok(())
    }

    #[test]
    fn new_security_policy() -> Result<(), Error> {
        use crate::security::Policy;

        let policy = Policy::from_version("default")?;
        config_builder(&policy)?;
        Ok(())
    }

    #[test]
    fn trust_location() -> Result<(), Error> {
        let pem_dir = Path::new(concat!(env!("CARGO_MANIFEST_DIR"), "/../../../tests/pems"));
        let mut cert = pem_dir.to_path_buf();
        cert.push("rsa_4096_sha512_client_cert.pem");
        let mut key = pem_dir.to_path_buf();
        key.push("rsa_4096_sha512_client_key.pem");

        let mut builder = crate::config::Builder::new();
        builder.set_security_policy(&security::DEFAULT_TLS13)?;
        builder.set_verify_host_callback(InsecureAcceptAllCertificatesHandler {})?;
        builder.load_pem(&fs::read(&cert)?, &fs::read(&key)?)?;
        builder.trust_location(Some(&cert), None)?;

        establish_connection(builder.build()?);
        Ok(())
    }

    /// `trust_location()` calls `s2n_config_set_verification_ca_location()`, which has the legacy behavior
    /// of enabling OCSP on clients. Since we do not want to replicate that behavior in the Rust bindings,
    /// this test verifies that `trust_location()` does not turn on OCSP. It also verifies that turning
    /// on OCSP explicitly still works when `trust_location()` is called.
    #[test]
    fn trust_location_does_not_change_ocsp_status() -> Result<(), Error> {
        let pem_dir = Path::new(concat!(env!("CARGO_MANIFEST_DIR"), "/../../../tests/pems"));
        let mut cert = pem_dir.to_path_buf();
        cert.push("rsa_4096_sha512_client_cert.pem");
        let mut key = pem_dir.to_path_buf();
        key.push("rsa_4096_sha512_client_key.pem");

        const OCSP_IANA_EXTENSION_ID: u16 = 5;

        for enable_ocsp in [true, false] {
            let config = {
                let mut config = crate::config::Builder::new();

                if enable_ocsp {
                    config.enable_ocsp()?;
                }

                config.set_security_policy(&security::DEFAULT_TLS13)?;
                config.set_verify_host_callback(InsecureAcceptAllCertificatesHandler {})?;
                config.set_client_hello_callback(HasExtensionClientHelloHandler {
                    // This client hello handler will check for the OCSP extension
                    extension_iana: OCSP_IANA_EXTENSION_ID,
                    extension_expected: enable_ocsp,
                })?;
                config.load_pem(&fs::read(&cert)?, &fs::read(&key)?)?;
                config.trust_location(Some(&cert), None)?;
                config.build()?
            };

            let mut pair = tls_pair(config);
            pair.server
                .0
                .connection_mut()
                .set_waker(Some(&noop_waker()))?;

            poll_tls_pair(pair);
        }
        Ok(())
    }

    #[test]
    fn connection_level_verify_host_callback() -> Result<(), Error> {
        let reject_config = {
            let keypair = CertKeyPair::default();
            let mut config = crate::config::Builder::new();
            // configure the config VerifyHostNameCallback to reject all certificates
            config.set_verify_host_callback(RejectAllCertificatesHandler {})?;
            config.set_security_policy(&security::DEFAULT_TLS13)?;
            config.load_pem(keypair.cert(), keypair.key())?;
            config.trust_pem(keypair.cert())?;
            config.set_client_auth_type(ClientAuthType::Required)?;
            config.build()?
        };

        // confirm that default connection establishment fails
        let mut pair = tls_pair(reject_config.clone());
        assert!(poll_tls_pair_result(&mut pair).is_err());

        // confirm that overriding the verify_host_callback on connection causes
        // the handshake to succeed
        pair = tls_pair(reject_config);
        pair.server
            .0
            .connection
            .set_verify_host_callback(InsecureAcceptAllCertificatesHandler {})
            .unwrap();
        pair.client
            .0
            .connection
            .set_verify_host_callback(InsecureAcceptAllCertificatesHandler {})
            .unwrap();
        assert!(poll_tls_pair_result(&mut pair).is_ok());

        Ok(())
    }

    #[test]
    fn no_client_auth() -> Result<(), Error> {
        use crate::enums::ClientAuthType;

        let config = {
            let mut config = config_builder(&security::DEFAULT_TLS13)?;
            config.set_client_auth_type(ClientAuthType::None)?;
            config.build()?
        };

        let server = {
            let mut server = crate::connection::Connection::new_server();
            server.set_config(config.clone())?;
            Harness::new(server)
        };

        let client = {
            let mut client = crate::connection::Connection::new_client();
            client.set_config(config)?;
            Harness::new(client)
        };

        let pair = Pair::new(server, client);
        let pair = poll_tls_pair(pair);
        let server = pair.server.0.connection;
        let client = pair.client.0.connection;

        for conn in [server, client] {
            assert!(!conn.client_cert_used());
            let cert = conn.client_cert_chain_bytes()?;
            assert!(cert.is_none());
            let sig_alg = conn.selected_client_signature_algorithm()?;
            assert!(sig_alg.is_none());
            let hash_alg = conn.selected_client_hash_algorithm()?;
            assert!(hash_alg.is_none());
        }

        Ok(())
    }

    #[test]
    fn client_auth() -> Result<(), Error> {
        use crate::enums::ClientAuthType;

        let config = {
            let mut config = config_builder(&security::DEFAULT_TLS13)?;
            config.set_client_auth_type(ClientAuthType::Optional)?;
            config.build()?
        };

        let server = {
            let mut server = crate::connection::Connection::new_server();
            server.set_config(config.clone())?;
            Harness::new(server)
        };

        let client = {
            let mut client = crate::connection::Connection::new_client();
            client.set_config(config)?;
            Harness::new(client)
        };

        let pair = Pair::new(server, client);
        let pair = poll_tls_pair(pair);
        let server = pair.server.0.connection;
        let client = pair.client.0.connection;

        let cert = server.client_cert_chain_bytes()?;
        assert!(cert.is_some());
        assert!(!cert.unwrap().is_empty());

        for conn in [server, client] {
            assert!(conn.client_cert_used());
            let sig_alg = conn.selected_client_signature_algorithm()?;
            assert!(sig_alg.is_some());
            let hash_alg = conn.selected_client_hash_algorithm()?;
            assert!(hash_alg.is_some());
        }

        Ok(())
    }

    #[test]
    fn system_certs_loaded_by_default() {
        let keypair = CertKeyPair::default();

        // Load the server certificate into the trust store by overriding the OpenSSL default
        // certificate location.
        temp_env::with_var("SSL_CERT_FILE", Some(keypair.cert_path()), || {
            let mut builder = Builder::new();
            builder
                .load_pem(keypair.cert(), keypair.key())
                .unwrap()
                .set_security_policy(&security::DEFAULT_TLS13)
                .unwrap()
                .set_verify_host_callback(InsecureAcceptAllCertificatesHandler {})
                .unwrap();

            let config = builder.build().unwrap();
            establish_connection(config);
        });
    }

    #[test]
    fn disable_loading_system_certs() {
        let keypair = CertKeyPair::default();

        // Load the server certificate into the trust store by overriding the OpenSSL default
        // certificate location.
        temp_env::with_var("SSL_CERT_FILE", Some(keypair.cert_path()), || {
            let mut builder = Builder::new();
            builder
                .load_pem(keypair.cert(), keypair.key())
                .unwrap()
                .set_security_policy(&security::DEFAULT_TLS13)
                .unwrap()
                .set_verify_host_callback(InsecureAcceptAllCertificatesHandler {})
                .unwrap();

            // Disable loading system certificates
            builder.with_system_certs(false).unwrap();

            let config = builder.build().unwrap();
            let mut config_with_system_certs = config.clone();

            let mut pair = tls_pair(config);

            // System certificates should not be loaded into the trust store. The handshake
            // should fail since the certificate should not be trusted.
            assert!(poll_tls_pair_result(&mut pair).is_err());

            // The handshake should succeed after trusting the certificate.
            unsafe {
                s2n_tls_sys::s2n_config_load_system_certs(config_with_system_certs.as_mut_ptr());
            }
            establish_connection(config_with_system_certs);
        });
    }
}


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

use std::path::{Path, PathBuf};

fn main() {
    let external = External::default();
    if external.is_enabled() {
        external.link();
    } else {
        #[cfg(feature = "cmake")]
        {
            // branch on a runtime value so we don't get unused code warnings
            if option_env("CARGO_FEATURE_CMAKE").is_some() {
                build_cmake();
            } else {
                build_vendored();
            }
        }

        #[cfg(not(feature = "cmake"))]
        build_vendored();
    }
}

fn env<N: AsRef<str>>(name: N) -> String {
    option_env(name).expect("missing env var")
}

fn option_env<N: AsRef<str>>(name: N) -> Option<String> {
    let name = name.as_ref();
    eprintln!("cargo:rerun-if-env-changed={}", name);
    std::env::var(name).ok()
}

struct FeatureDetector<'a> {
    builder: cc::Build,
    out_dir: &'a Path,
}

impl<'a> FeatureDetector<'a> {
    pub fn new(out_dir: &'a Path) -> Self {
        let builder = builder();
        Self { builder, out_dir }
    }

    pub fn supports(&self, name: &str) -> bool {
        let mut build = self.builder.get_compiler().to_command();

        let global_flags = std::path::Path::new("lib/tests/features/GLOBAL.flags");
        assert!(
            global_flags.exists(),
            "missing flags file: {:?}",
            global_flags.display()
        );

        let global_flags = std::fs::read_to_string(global_flags).unwrap();
        for flag in global_flags.trim().split(' ').filter(|f| !f.is_empty()) {
            build.arg(flag);
        }

        let base = std::path::Path::new("lib/tests/features").join(name);

        let file = base.with_extension("c");
        assert!(file.exists(), "missing feature file: {:?}", file.display());

        let probe_flags = base.with_extension("flags");
        assert!(
            probe_flags.exists(),
            "missing flags file: {:?}",
            probe_flags.display()
        );

        let probe_flags = std::fs::read_to_string(probe_flags).unwrap();
        for flag in probe_flags.trim().split(' ').filter(|f| !f.is_empty()) {
            build.arg(flag);
        }

        build
            // just compile the file and don't link
            .arg("-c")
            .arg("-o")
            .arg(self.out_dir.join(name).with_extension("o"))
            .arg(&file);

        eprintln!("=== Testing feature {name} ===");
        build.status().unwrap().success()
    }
}

fn build_vendored() {
    let mut build = builder();

    let pq = option_env("CARGO_FEATURE_PQ").is_some();

    // TODO each pq section needs to be built separately since it
    //      has its own relative include paths
    assert!(!pq, "pq builds are not currently supported without cmake");

    build.files(include!("./files.rs").iter().copied().filter(|file| {
        // the pq entry file is still needed
        if *file == "lib/pq-crypto/s2n_pq.c" {
            return true;
        }

        if file.starts_with("lib/pq-crypto/") {
            return pq;
        }

        true
    }));

    if env("PROFILE") == "release" {
        // fortify source is only available in release mode
        build.define("_FORTIFY_SOURCE", "2");
        build.define("NDEBUG", "1");

        // build s2n-tls with LTO if supported
        if build.get_compiler().is_like_gnu() {
            build
                .flag_if_supported("-flto")
                .flag_if_supported("-ffat-lto-objects");
        }
    }

    if !pq {
        build.define("S2N_NO_PQ", "1");
    }

    let out_dir = PathBuf::from(env("OUT_DIR"));

    let features = FeatureDetector::new(&out_dir);

    let mut feature_names = std::fs::read_dir("lib/tests/features")
        .expect("missing features directory")
        .flatten()
        .filter(|file| {
            let file = file.path();
            file.extension().map_or(false, |ext| ext == "c")
        })
        .map(|file| {
            file.path()
                .file_stem()
                .unwrap()
                .to_str()
                .unwrap()
                .to_string()
        })
        .collect::<Vec<_>>();

    feature_names.sort();

    for name in &feature_names {
        let is_supported = features.supports(name);
        eprintln!("{name}: {is_supported}");
        if is_supported {
            build.define(name, "1");

            // stacktraces are only available if execinfo is
            if name == "S2N_EXECINFO_AVAILABLE" && option_env("CARGO_FEATURE_STACKTRACE").is_some()
            {
                build.define("S2N_STACKTRACE", "1");
            }
        }
    }

    // don't spit out a bunch of warnings to the end user, since they won't really be able
    // to do anything with it
    build.warnings(false);

    build.compile("s2n-tls");

    // tell rust we're linking with libcrypto
    println!("cargo:rustc-link-lib=crypto");

    // let consumers know where to find our header files
    let include_dir = out_dir.join("include");
    std::fs::create_dir_all(&include_dir).unwrap();
    std::fs::copy("lib/api/s2n.h", include_dir.join("s2n.h")).unwrap();
    println!("cargo:include={}", include_dir.display());
}

fn builder() -> cc::Build {
    let mut build = cc::Build::new();

    build
        // pull the include path from the openssl-sys dependency
        .include(env("DEP_OPENSSL_INCLUDE"))
        .include("lib")
        .include("lib/api")
        .flag("-std=c11")
        .flag("-fgnu89-inline")
        // make sure the stack is non-executable
        .flag_if_supported("-z relro")
        .flag_if_supported("-z now")
        .flag_if_supported("-z noexecstack")
        // we use some deprecated libcrypto features so don't warn here
        .flag_if_supported("-Wno-deprecated-declarations")
        .define("_POSIX_C_SOURCE", "200112L");

    build
}

#[cfg(feature = "cmake")]
fn build_cmake() {
    let mut config = cmake::Config::new("lib");

    // sometimes openssl-sys decides not to set this value so we may need to set it anyway
    if option_env("DEP_OPENSSL_ROOT").is_none() {
        let include = env("DEP_OPENSSL_INCLUDE");
        if let Some(root) = Path::new(&include).parent() {
            std::env::set_var("DEP_OPENSSL_ROOT", root);
        }
    }

    config
        .register_dep("openssl")
        .configure_arg("-DBUILD_TESTING=off");

    if option_env("CARGO_FEATURE_PQ").is_none() {
        config.configure_arg("-DS2N_NO_PQ=on");
    }

    let dst = config.build();

    let lib = search(dst.join("lib64"))
        .or_else(|| search(dst.join("lib")))
        .or_else(|| search(dst.join("build").join("lib")))
        .expect("could not build libs2n");

    // link the built artifact
    if lib.join("libs2n.a").exists() {
        println!("cargo:rustc-link-lib=static=s2n");
    } else {
        println!("cargo:rustc-link-lib=s2n");
    }

    println!("cargo:include={}", dst.join("include").display());

    // tell rust we're linking with libcrypto
    println!("cargo:rustc-link-lib=crypto");

    fn search(path: PathBuf) -> Option<PathBuf> {
        if path.exists() {
            println!("cargo:rustc-link-search={}", path.display());
            Some(path)
        } else {
            None
        }
    }
}

struct External {
    lib_dir: Option<PathBuf>,
    include_dir: Option<PathBuf>,
}

impl Default for External {
    fn default() -> Self {
        let dir = option_env("S2N_TLS_DIR").map(PathBuf::from);

        let lib_dir = option_env("S2N_TLS_LIB_DIR")
            .map(PathBuf::from)
            .or_else(|| dir.as_ref().map(|d| d.join("lib")));

        let include_dir = option_env("S2N_TLS_INCLUDE_DIR")
            .map(PathBuf::from)
            .or_else(|| dir.as_ref().map(|d| d.join("include")));

        Self {
            lib_dir,
            include_dir,
        }
    }
}

impl External {
    fn is_enabled(&self) -> bool {
        self.lib_dir.is_some()
    }

    fn link(&self) {
        println!(
            "cargo:rustc-link-search={}",
            self.lib_dir.as_ref().unwrap().display()
        );
        println!("cargo:rustc-link-lib=s2n");

        // tell rust we're linking with libcrypto
        println!("cargo:rustc-link-lib=crypto");

        if let Some(include_dir) = self.include_dir.as_ref() {
            println!("cargo:include={}", include_dir.display());
        }
    }
}


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

#[rustfmt::skip]
mod api;

pub use api::*;

mod features;
pub use features::*;

// Additional defines that don't get imported with bindgen

pub mod s2n_status_code {
    pub type Type = libc::c_int;
    pub const SUCCESS: Type = 0;
    pub const FAILURE: Type = -1;
}

pub mod s2n_tls_version {
    pub type Type = libc::c_int;
    pub const SSLV2: Type = 20;
    pub const SSLV3: Type = 30;
    pub const TLS10: Type = 31;
    pub const TLS11: Type = 32;
    pub const TLS12: Type = 33;
    pub const TLS13: Type = 34;
    pub const UNKNOWN: Type = 0;
}

#[cfg(test)]
#[rustfmt::skip]
mod tests;


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

use s2n_tls_sys::*;

#[test]
fn s2n_init_test() {
    unsafe {
        // don't force the tests to use mlock
        std::env::set_var("S2N_DONT_MLOCK", "1");

        // try to initialize the library
        if s2n_init() != 0 {
            let error = *s2n_errno_location();
            let msg = s2n_strerror_name(error);
            let msg = std::ffi::CStr::from_ptr(msg);
            panic!("s2n did not initialize correctly: {:?}", msg);
        }
    }
}


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

use clap::Parser;
use s2n_tls::{config::Config, security::DEFAULT_TLS13};
use s2n_tls_tokio::TlsConnector;
use std::{error::Error, fs};
use tokio::{io::AsyncWriteExt, net::TcpStream};

/// NOTE: this certificate is to be used for demonstration purposes only!
const DEFAULT_CERT: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/examples/certs/cert.pem");

#[derive(Parser, Debug)]
struct Args {
    #[clap(short, long, default_value_t = String::from(DEFAULT_CERT))]
    trust: String,
    addr: String,
}

async fn run_client(trust_pem: &[u8], addr: &str) -> Result<(), Box<dyn Error>> {
    // Set up the configuration for new connections.
    // Minimally you will need a trust store.
    let mut config = Config::builder();
    config.set_security_policy(&DEFAULT_TLS13)?;
    config.trust_pem(trust_pem)?;

    // Create the TlsConnector based on the configuration.
    let client = TlsConnector::new(config.build()?);

    // Connect to the server.
    let stream = TcpStream::connect(addr).await?;
    let tls = client.connect("localhost", stream).await?;
    println!("{:#?}", tls);

    // Split the stream.
    // This allows us to call read and write from different tasks.
    let (mut reader, mut writer) = tokio::io::split(tls);

    // Copy data from the server to stdout
    tokio::spawn(async move {
        let mut stdout = tokio::io::stdout();
        tokio::io::copy(&mut reader, &mut stdout).await
    });

    // Send data from stdin to the server
    let mut stdin = tokio::io::stdin();
    tokio::io::copy(&mut stdin, &mut writer).await?;
    writer.shutdown().await?;

    Ok(())
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let args = Args::parse();
    let trust_pem = fs::read(args.trust)?;
    run_client(&trust_pem, &args.addr).await?;
    Ok(())
}


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

use clap::Parser;
use s2n_tls::{config::Config, enums::Mode, pool::ConfigPoolBuilder, security::DEFAULT_TLS13};
use s2n_tls_tokio::TlsAcceptor;
use std::{error::Error, fs};
use tokio::{io::AsyncWriteExt, net::TcpListener};

/// NOTE: this certificate and key are to be used for demonstration purposes only!
const DEFAULT_CERT: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/examples/certs/cert.pem");
const DEFAULT_KEY: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/examples/certs/key.pem");

#[derive(Parser, Debug)]
struct Args {
    #[clap(short, long, requires = "key", default_value_t = String::from(DEFAULT_CERT))]
    cert: String,
    #[clap(short, long, requires = "cert", default_value_t = String::from(DEFAULT_KEY))]
    key: String,
    #[clap(short, long, default_value_t = String::from("127.0.0.1:0"))]
    addr: String,
}

async fn run_server(cert_pem: &[u8], key_pem: &[u8], addr: &str) -> Result<(), Box<dyn Error>> {
    // Set up the configuration for new connections.
    // Minimally you will need a certificate and private key.
    let mut config = Config::builder();
    config.set_security_policy(&DEFAULT_TLS13)?;
    config.load_pem(cert_pem, key_pem)?;

    // Create a connection pool to reuse connections.
    let mut pool = ConfigPoolBuilder::new(Mode::Server, config.build()?);
    pool.set_max_pool_size(10);

    // Create the TlsAcceptor based on the pool.
    let server = TlsAcceptor::new(pool.build());

    // Bind to an address and listen for connections.
    // ":0" can be used to automatically assign a port.
    let listener = TcpListener::bind(&addr).await?;
    let addr = listener
        .local_addr()
        .map(|x| x.to_string())
        .unwrap_or_else(|_| "UNKNOWN".to_owned());
    println!("Listening on {}", addr);

    loop {
        // Wait for a client to connect.
        let (stream, peer_addr) = listener.accept().await?;
        println!("Connection from {:?}", peer_addr);

        // Spawn a new task to handle the connection.
        // We probably want to spawn the task BEFORE calling TcpAcceptor::accept,
        // because the TLS handshake can be slow.
        let server = server.clone();
        tokio::spawn(async move {
            let mut tls = server.accept(stream).await?;
            println!("{:#?}", tls);

            // Copy data from the client to stdout
            let mut stdout = tokio::io::stdout();
            tokio::io::copy(&mut tls, &mut stdout).await?;
            tls.shutdown().await?;
            println!("Connection from {:?} closed", peer_addr);

            Ok::<(), Box<dyn Error + Send + Sync>>(())
        });
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let args = Args::parse();
    let cert_pem = fs::read(args.cert)?;
    let key_pem = fs::read(args.key)?;
    run_server(&cert_pem, &key_pem, &args.addr).await?;
    Ok(())
}


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

use errno::{set_errno, Errno};
use pin_project_lite::pin_project;
use s2n_tls::{
    config::Config,
    connection::{Builder, Connection},
    enums::{Blinding, CallbackResult, Mode},
    error::Error,
};
use std::{
    fmt,
    future::Future,
    io, mem,
    os::raw::{c_int, c_void},
    pin::Pin,
    task::{
        Context, Poll,
        Poll::{Pending, Ready},
    },
};
use tokio::{
    io::{AsyncRead, AsyncWrite, ReadBuf},
    time::{sleep, Duration, Sleep},
};

macro_rules! ready {
    ($x:expr) => {
        match $x {
            Ready(r) => r,
            Pending => return Pending,
        }
    };
}

#[derive(Clone)]
pub struct TlsAcceptor<B: Builder = Config>
where
    <B as Builder>::Output: Unpin,
{
    builder: B,
}

impl<B: Builder> TlsAcceptor<B>
where
    <B as Builder>::Output: Unpin,
{
    pub fn new(builder: B) -> Self {
        TlsAcceptor { builder }
    }

    pub async fn accept<S>(&self, stream: S) -> Result<TlsStream<S, B::Output>, Error>
    where
        S: AsyncRead + AsyncWrite + Unpin,
    {
        let conn = self.builder.build_connection(Mode::Server)?;
        TlsStream::open(conn, stream).await
    }
}

#[derive(Clone)]
pub struct TlsConnector<B: Builder = Config>
where
    <B as Builder>::Output: Unpin,
{
    builder: B,
}

impl<B: Builder> TlsConnector<B>
where
    <B as Builder>::Output: Unpin,
{
    pub fn new(builder: B) -> Self {
        TlsConnector { builder }
    }

    pub async fn connect<S>(
        &self,
        domain: &str,
        stream: S,
    ) -> Result<TlsStream<S, B::Output>, Error>
    where
        S: AsyncRead + AsyncWrite + Unpin,
    {
        let mut conn = self.builder.build_connection(Mode::Client)?;
        conn.as_mut().set_server_name(domain)?;
        TlsStream::open(conn, stream).await
    }
}

struct TlsHandshake<'a, S, C>
where
    C: AsRef<Connection> + AsMut<Connection> + Unpin,
    S: AsyncRead + AsyncWrite + Unpin,
{
    tls: &'a mut TlsStream<S, C>,
    error: Option<Error>,
}

impl<S, C> Future for TlsHandshake<'_, S, C>
where
    C: AsRef<Connection> + AsMut<Connection> + Unpin,
    S: AsyncRead + AsyncWrite + Unpin,
{
    type Output = Result<(), Error>;

    fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
        // Retrieve a result, either from the stored error
        // or by polling Connection::poll_negotiate().
        // Connection::poll_negotiate() only completes once,
        // regardless of how often this method is polled.
        let result = match self.error.take() {
            Some(err) => Err(err),
            None => {
                ready!(self.tls.with_io(ctx, |context| {
                    let conn = context.get_mut().as_mut();
                    conn.poll_negotiate().map(|r| r.map(|_| ()))
                }))
            }
        };
        // If the result isn't a fatal error, return it immediately.
        // Otherwise, poll Connection::poll_shutdown().
        //
        // Shutdown is only best-effort.
        // When Connection::poll_shutdown() completes, even with an error,
        // we return the original Connection::poll_negotiate() error.
        match result {
            Ok(r) => Ok(r).into(),
            Err(e) if e.is_retryable() => Err(e).into(),
            Err(e) => match Pin::new(&mut self.tls).poll_shutdown(ctx) {
                Pending => {
                    self.error = Some(e);
                    Pending
                }
                Ready(_) => Err(e).into(),
            },
        }
    }
}

pin_project! {
struct BlindingState {
    #[pin]
    timer: Sleep,

    // The remembered error if we got into blinding because of
    // an error, or Ok(()) if we didn't. After returning the error,
    // this goes back to Ok(()).
    remembered_error: Result<(), Error>,
}
}

pub struct TlsStream<S, C = Connection>
where
    C: AsRef<Connection> + AsMut<Connection> + Unpin,
    S: AsyncRead + AsyncWrite + Unpin,
{
    conn: C,
    stream: S,
    blinding: Option<Pin<Box<BlindingState>>>,
}

impl<S, C> TlsStream<S, C>
where
    C: AsRef<Connection> + AsMut<Connection> + Unpin,
    S: AsyncRead + AsyncWrite + Unpin,
{
    ///Access a shared reference to the underlaying io stream
    pub fn get_ref(&self) -> &S {
        &self.stream
    }

    ///Access the mutable reference to the underlaying io stream
    pub fn get_mut(&mut self) -> &mut S {
        &mut self.stream
    }

    async fn open(mut conn: C, stream: S) -> Result<Self, Error> {
        conn.as_mut().set_blinding(Blinding::SelfService)?;
        let mut tls = TlsStream {
            conn,
            stream,
            blinding: None,
        };
        TlsHandshake {
            tls: &mut tls,
            error: None,
        }
        .await?;
        Ok(tls)
    }

    fn with_io<F, R>(&mut self, ctx: &mut Context, action: F) -> Poll<Result<R, Error>>
    where
        F: FnOnce(Pin<&mut Self>) -> Poll<Result<R, Error>>,
    {
        // Setting contexts on a connection is considered unsafe
        // because the raw pointers provide no lifetime or memory guarantees.
        // We protect against this by pinning the stream during the action
        // and clearing the context afterwards.
        unsafe {
            let context = self as *mut Self as *mut c_void;

            self.as_mut().set_receive_callback(Some(Self::recv_io_cb))?;
            self.as_mut().set_send_callback(Some(Self::send_io_cb))?;
            self.as_mut().set_receive_context(context)?;
            self.as_mut().set_send_context(context)?;
            self.as_mut().set_waker(Some(ctx.waker()))?;

            let result = action(Pin::new(self));

            self.as_mut().set_receive_callback(None)?;
            self.as_mut().set_send_callback(None)?;
            self.as_mut().set_receive_context(std::ptr::null_mut())?;
            self.as_mut().set_send_context(std::ptr::null_mut())?;
            self.as_mut().set_waker(None)?;
            result
        }
    }

    fn poll_io<F>(ctx: *mut c_void, action: F) -> c_int
    where
        F: FnOnce(Pin<&mut S>, &mut Context) -> Poll<Result<usize, std::io::Error>>,
    {
        debug_assert_ne!(ctx, std::ptr::null_mut());
        let tls = unsafe { &mut *(ctx as *mut Self) };

        let mut async_context = Context::from_waker(tls.conn.as_ref().waker().unwrap());
        let stream = Pin::new(&mut tls.stream);

        match action(stream, &mut async_context) {
            Poll::Ready(Ok(len)) => len as c_int,
            Poll::Pending => {
                set_errno(Errno(libc::EWOULDBLOCK));
                CallbackResult::Failure.into()
            }
            _ => CallbackResult::Failure.into(),
        }
    }

    unsafe extern "C" fn recv_io_cb(ctx: *mut c_void, buf: *mut u8, len: u32) -> c_int {
        Self::poll_io(ctx, |stream, async_context| {
            let mut dest = ReadBuf::new(std::slice::from_raw_parts_mut(buf, len as usize));
            stream
                .poll_read(async_context, &mut dest)
                .map_ok(|_| dest.filled().len())
        })
    }

    unsafe extern "C" fn send_io_cb(ctx: *mut c_void, buf: *const u8, len: u32) -> c_int {
        Self::poll_io(ctx, |stream, async_context| {
            let src = std::slice::from_raw_parts(buf, len as usize);
            stream.poll_write(async_context, src)
        })
    }

    // Sets the blinding timer to the remaining blinding delay and possibly
    // remembers an error.
    //
    // Returns the error if there was no blinding needed and the error
    // did not need to be remembered.
    fn set_blinding_timer(
        self: Pin<&mut Self>,
        mut remembered_error: Result<(), Error>,
    ) -> Result<(), Error> {
        let mut tls = self.get_mut();

        if tls.blinding.is_none() {
            let delay = tls.as_ref().remaining_blinding_delay()?;
            if !delay.is_zero() {
                // Sleep operates at the milisecond resolution, so add an extra
                // millisecond to account for any stray nanoseconds.
                let safety = Duration::from_millis(1);
                // Return the error *later*, after the blinding is done
                let remembered_error = mem::replace(&mut remembered_error, Ok(()));
                tls.blinding = Some(Box::pin(BlindingState {
                    timer: sleep(delay.saturating_add(safety)),
                    remembered_error,
                }));
            }
        }

        remembered_error
    }

    /// Polls the blinding timer, if there is any.
    ///
    /// s2n has a "blinding" functionality - when a bad behavior from the peer
    /// is detected, sleeps for 10-30 seconds before answering the client
    /// and closing the connection. This mitigates some timing side channels
    /// that could leak information about encrypted data. See the
    /// `s2n_connection_set_blinding` docs for more details.
    ///
    /// For security reasons, to allow for blinding to correctly function,
    /// before dropping an s2n connection, you should wait until either
    /// `poll_blinding` or `poll_shutdown` (which calls `poll_blinding`
    /// internally) returns ready.
    pub fn poll_blinding(
        mut self: Pin<&mut Self>,
        ctx: &mut Context<'_>,
    ) -> Poll<Result<(), Error>> {
        self.as_mut().set_blinding_timer(Ok(()))?;

        let tls = self.get_mut();

        if let Some(blinding) = &mut tls.blinding {
            ready!(blinding.as_mut().project().timer.as_mut().poll(ctx));

            // Set blinding to None to ensure the next go can have blinding
            let mut blinding = tls.blinding.take().unwrap();

            // If there is an error, return it
            mem::replace(blinding.as_mut().project().remembered_error, Ok(()))?;
        }

        // Otherwise we are OK
        Poll::Ready(Ok(()))
    }

    pub async fn apply_blinding(&mut self) -> Result<(), Error> {
        ApplyBlinding { stream: self }.await
    }
}

impl<S, C> AsRef<Connection> for TlsStream<S, C>
where
    C: AsRef<Connection> + AsMut<Connection> + Unpin,
    S: AsyncRead + AsyncWrite + Unpin,
{
    fn as_ref(&self) -> &Connection {
        self.conn.as_ref()
    }
}

impl<S, C> AsMut<Connection> for TlsStream<S, C>
where
    C: AsRef<Connection> + AsMut<Connection> + Unpin,
    S: AsyncRead + AsyncWrite + Unpin,
{
    fn as_mut(&mut self) -> &mut Connection {
        self.conn.as_mut()
    }
}

impl<S, C> AsyncRead for TlsStream<S, C>
where
    C: AsRef<Connection> + AsMut<Connection> + Unpin,
    S: AsyncRead + AsyncWrite + Unpin,
{
    fn poll_read(
        self: Pin<&mut Self>,
        ctx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        let tls = self.get_mut();
        tls.with_io(ctx, |mut context| {
            context
                .conn
                .as_mut()
                // Safe since poll_recv_uninitialized does not
                // deinitialize any bytes.
                .poll_recv_uninitialized(unsafe { buf.unfilled_mut() })
                .map_ok(|size| {
                    unsafe {
                        // Safe since poll_recv_uninitialized guaranteed
                        // us that the first `size` bytes have been
                        // initialized.
                        buf.assume_init(size);
                    }
                    buf.advance(size);
                })
        })
        .map_err(io::Error::from)
    }
}

impl<S, C> AsyncWrite for TlsStream<S, C>
where
    C: AsRef<Connection> + AsMut<Connection> + Unpin,
    S: AsyncRead + AsyncWrite + Unpin,
{
    fn poll_write(
        self: Pin<&mut Self>,
        ctx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        let tls = self.get_mut();
        tls.with_io(ctx, |mut context| context.conn.as_mut().poll_send(buf))
            .map_err(io::Error::from)
    }

    fn poll_flush(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<io::Result<()>> {
        let tls = self.get_mut();

        ready!(tls.with_io(ctx, |mut context| {
            context.conn.as_mut().poll_flush().map(|r| r.map(|_| ()))
        }))
        .map_err(io::Error::from)?;

        Pin::new(&mut tls.stream).poll_flush(ctx)
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<io::Result<()>> {
        ready!(self.as_mut().poll_blinding(ctx))?;

        let status = ready!(self.as_mut().with_io(ctx, |mut context| {
            context.conn.as_mut().poll_shutdown().map(|r| r.map(|_| ()))
        }));

        if let Err(e) = status {
            // In case of an error shutting down, make sure you wait for
            // the blinding timeout.
            self.as_mut().set_blinding_timer(Err(e))?;
            ready!(self.as_mut().poll_blinding(ctx))?;
            unreachable!("should have returned the error we just put in!");
        }

        Pin::new(&mut self.as_mut().stream).poll_shutdown(ctx)
    }
}

impl<S, C> fmt::Debug for TlsStream<S, C>
where
    C: AsRef<Connection> + AsMut<Connection> + Unpin,
    S: AsyncRead + AsyncWrite + Unpin,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("TlsStream")
            .field("connection", self.as_ref())
            .finish()
    }
}

struct ApplyBlinding<'a, S, C>
where
    C: AsRef<Connection> + AsMut<Connection> + Unpin,
    S: AsyncRead + AsyncWrite + Unpin,
{
    stream: &'a mut TlsStream<S, C>,
}

impl<'a, S, C> Future for ApplyBlinding<'a, S, C>
where
    C: AsRef<Connection> + AsMut<Connection> + Unpin,
    S: AsyncRead + AsyncWrite + Unpin,
{
    type Output = Result<(), Error>;

    fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
        Pin::new(&mut *self.as_mut().stream).poll_blinding(ctx)
    }
}


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

use rand::Rng;
use s2n_tls::{
    config::Config,
    connection::{Connection, ModifiedBuilder},
    enums::{ClientAuthType, Mode, Version},
    error::{Error, ErrorType},
    pool::ConfigPoolBuilder,
    security::DEFAULT_TLS13,
};
use s2n_tls_tokio::{TlsAcceptor, TlsConnector};
use std::{collections::VecDeque, time::Duration};
use tokio::time;

pub mod common;

#[tokio::test]
async fn handshake_basic() -> Result<(), Box<dyn std::error::Error>> {
    let (server_stream, client_stream) = common::get_streams().await?;

    let client = TlsConnector::new(common::client_config()?.build()?);
    let server = TlsAcceptor::new(common::server_config()?.build()?);

    let (client_result, server_result) =
        common::run_negotiate(&client, client_stream, &server, server_stream).await?;

    for tls in [client_result, server_result] {
        // Security policy ensures TLS1.3.
        assert_eq!(tls.as_ref().actual_protocol_version()?, Version::TLS13);
        // Handshake types may change, but will at least be negotiated.
        assert!(tls.as_ref().handshake_type()?.contains("NEGOTIATED"));
        // Cipher suite may change, so just makes sure we can retrieve it.
        assert!(tls.as_ref().cipher_suite().is_ok());
        assert!(tls.as_ref().selected_curve().is_ok());
    }

    Ok(())
}

#[tokio::test(flavor = "multi_thread")]
async fn handshake_with_pool_multithread() -> Result<(), Box<dyn std::error::Error>> {
    const COUNT: usize = 20;
    const CLIENT_LIMIT: usize = 3;

    let client_config = common::client_config()?.build()?;
    let server_config = common::server_config()?.build()?;

    let mut client_pool = ConfigPoolBuilder::new(Mode::Client, client_config);
    client_pool.set_max_pool_size(CLIENT_LIMIT);

    let client_pool = client_pool.build();
    let server_pool = ConfigPoolBuilder::new(Mode::Server, server_config).build();

    let client = TlsConnector::new(client_pool.clone());
    let server = TlsAcceptor::new(server_pool.clone());

    let mut tasks = VecDeque::new();
    for _ in 0..COUNT {
        let client = client.clone();
        let server = server.clone();
        tasks.push_back(tokio::spawn(async move {
            // Start each handshake at a randomly determined time
            let rand = rand::thread_rng().gen_range(0..50);
            time::sleep(Duration::from_millis(rand)).await;

            let (server_stream, client_stream) = common::get_streams().await.unwrap();
            common::run_negotiate(&client, client_stream, &server, server_stream).await
        }));
    }

    for task in tasks {
        task.await??;
    }
    Ok(())
}

#[tokio::test]
async fn handshake_with_connection_config() -> Result<(), Box<dyn std::error::Error>> {
    // Setup the client with a method
    fn with_client_auth(conn: &mut Connection) -> Result<&mut Connection, Error> {
        conn.set_client_auth_type(ClientAuthType::Optional)
    }
    let client_builder = ModifiedBuilder::new(common::client_config()?.build()?, with_client_auth);

    // Setup the server with a closure
    let server_builder = ModifiedBuilder::new(common::server_config()?.build()?, |conn| {
        conn.set_client_auth_type(ClientAuthType::Optional)
    });

    let client = TlsConnector::new(client_builder);
    let server = TlsAcceptor::new(server_builder);

    let (server_stream, client_stream) = common::get_streams().await?;
    let (client_result, server_result) =
        common::run_negotiate(&client, client_stream, &server, server_stream).await?;

    for tls in [client_result, server_result] {
        assert!(tls.as_ref().handshake_type()?.contains("CLIENT_AUTH"));
    }

    Ok(())
}

#[tokio::test]
async fn handshake_with_connection_config_with_pool() -> Result<(), Box<dyn std::error::Error>> {
    fn with_client_auth(conn: &mut Connection) -> Result<&mut Connection, Error> {
        conn.set_client_auth_type(ClientAuthType::Optional)
    }
    let client_builder = ModifiedBuilder::new(common::client_config()?.build()?, with_client_auth);
    let server_pool =
        ConfigPoolBuilder::new(Mode::Server, common::server_config()?.build()?).build();
    let server_builder = ModifiedBuilder::new(server_pool, with_client_auth);

    let client = TlsConnector::new(client_builder);
    let server = TlsAcceptor::new(server_builder);

    for _ in 0..5 {
        let (server_stream, client_stream) = common::get_streams().await?;
        let (_, server_result) =
            common::run_negotiate(&client, client_stream, &server, server_stream).await?;
        assert!(server_result
            .as_ref()
            .handshake_type()?
            .contains("CLIENT_AUTH"));
    }

    Ok(())
}

#[tokio::test]
async fn handshake_error() -> Result<(), Box<dyn std::error::Error>> {
    // Config::default() does not include any RSA certificates,
    // but only provides TLS1.2 cipher suites that require RSA auth.
    // The server will fail to choose a cipher suite, but
    // S2N_ERR_CIPHER_NOT_SUPPORTED is specifically excluded from blinding.
    let bad_config = Config::default();
    let client_config = common::client_config()?.build()?;
    let server_config = bad_config;

    let client = TlsConnector::new(client_config);
    let server = TlsAcceptor::new(server_config);

    let (server_stream, client_stream) = common::get_streams().await?;
    let result = common::run_negotiate(&client, client_stream, &server, server_stream).await;
    assert!(matches!(result, Err(e) if !e.is_retryable()));

    Ok(())
}

#[tokio::test(start_paused = true)]
async fn handshake_error_with_blinding() -> Result<(), Box<dyn std::error::Error>> {
    let clock = common::TokioTime::default();

    // Config::builder() does not include a trust store.
    // The client will reject the server certificate as untrusted.
    let mut bad_config = Config::builder();
    bad_config.set_security_policy(&DEFAULT_TLS13)?;
    bad_config.set_monotonic_clock(clock)?;
    let client_config = bad_config.build()?;
    let server_config = common::server_config()?.build()?;

    let client = TlsConnector::new(client_config.clone());
    let server = TlsAcceptor::new(server_config.clone());

    // Handshake MUST NOT finish faster than minimal blinding time.
    let (server_stream, client_stream) = common::get_streams().await?;
    let timeout = time::timeout(
        common::MIN_BLINDING_SECS,
        common::run_negotiate(&client, client_stream, &server, server_stream),
    )
    .await;
    assert!(timeout.is_err());

    // Handshake MUST eventually gracefully close after blinding
    let (server_stream, client_stream) = common::get_streams().await?;
    let timeout = time::timeout(
        common::MAX_BLINDING_SECS.mul_f32(1.1),
        common::run_negotiate(&client, client_stream, &server, server_stream),
    )
    .await;
    let result = timeout?;
    assert!(result.is_err());
    assert_eq!(result.unwrap_err().kind(), ErrorType::ProtocolError);

    Ok(())
}

#[tokio::test]
async fn io_stream_access() -> Result<(), Box<dyn std::error::Error>> {
    let (server_stream, client_stream) = common::get_streams().await?;

    let client_addr = client_stream.local_addr().unwrap();
    let client = TlsConnector::new(common::client_config()?.build()?);
    let server = TlsAcceptor::new(common::server_config()?.build()?);

    let (mut client_result, _server_result) =
        common::run_negotiate(&client, client_stream, &server, server_stream).await?;

    assert_eq!(client_result.get_ref().local_addr().unwrap(), client_addr);
    assert_eq!(client_result.get_mut().local_addr().unwrap(), client_addr);

    Ok(())
}


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

use s2n_tls_tokio::{TlsAcceptor, TlsConnector};
use std::{io, task::Poll::*};
use tokio::io::{AsyncReadExt, AsyncWriteExt};

pub mod common;

const TEST_DATA: &[u8] = "hello world".as_bytes();

// The maximum TLS record payload is 2^14 bytes.
// Send more to ensure multiple records.
const LARGE_TEST_DATA: &[u8] = &[5; (1 << 15)];

#[tokio::test]
async fn send_and_recv_basic() -> Result<(), Box<dyn std::error::Error>> {
    let (server_stream, client_stream) = common::get_streams().await?;

    let connector = TlsConnector::new(common::client_config()?.build()?);
    let acceptor = TlsAcceptor::new(common::server_config()?.build()?);

    let (mut client, mut server) =
        common::run_negotiate(&connector, client_stream, &acceptor, server_stream).await?;

    client.write_all(TEST_DATA).await?;

    let mut received = [0; TEST_DATA.len()];
    assert_eq!(server.read_exact(&mut received).await?, TEST_DATA.len());
    assert_eq!(TEST_DATA, received);

    Ok(())
}

#[tokio::test]
async fn send_and_recv_into_vec() -> Result<(), Box<dyn std::error::Error>> {
    let (server_stream, client_stream) = common::get_streams().await?;

    let connector = TlsConnector::new(common::client_config()?.build()?);
    let acceptor = TlsAcceptor::new(common::server_config()?.build()?);

    let (mut client, mut server) =
        common::run_negotiate(&connector, client_stream, &acceptor, server_stream).await?;

    client.write_all(TEST_DATA).await?;

    let mut received = vec![];
    while received.len() < TEST_DATA.len() {
        let bytes_read = server.read_buf(&mut received).await?;
        assert!(bytes_read > 0);
    }
    assert_eq!(TEST_DATA, received);

    Ok(())
}

#[tokio::test]
async fn send_and_recv_multiple_records() -> Result<(), Box<dyn std::error::Error>> {
    let (server_stream, client_stream) = common::get_streams().await?;

    let connector = TlsConnector::new(common::client_config()?.build()?);
    let acceptor = TlsAcceptor::new(common::server_config()?.build()?);

    let (mut client, mut server) =
        common::run_negotiate(&connector, client_stream, &acceptor, server_stream).await?;

    let mut received = [0; LARGE_TEST_DATA.len()];
    let (_, read_size) = tokio::try_join!(
        client.write_all(LARGE_TEST_DATA),
        server.read_exact(&mut received)
    )?;
    assert_eq!(LARGE_TEST_DATA.len(), read_size);
    assert_eq!(LARGE_TEST_DATA, received);

    Ok(())
}

#[tokio::test]
async fn send_and_recv_split() -> Result<(), Box<dyn std::error::Error>> {
    let (server_stream, client_stream) = common::get_streams().await?;

    let connector = TlsConnector::new(common::client_config()?.build()?);
    let acceptor = TlsAcceptor::new(common::server_config()?.build()?);

    let (client, server) =
        common::run_negotiate(&connector, client_stream, &acceptor, server_stream).await?;

    let (mut client_read, mut client_write) = tokio::io::split(client);
    let (mut server_read, mut server_write) = tokio::io::split(server);

    let mut client_received = [0; LARGE_TEST_DATA.len()];
    let mut server_received = [0; LARGE_TEST_DATA.len()];
    let (_, _, client_bytes, server_bytes) = tokio::try_join!(
        client_write.write_all(LARGE_TEST_DATA),
        server_write.write_all(LARGE_TEST_DATA),
        client_read.read_exact(&mut client_received),
        server_read.read_exact(&mut server_received)
    )?;

    assert_eq!(client_bytes, LARGE_TEST_DATA.len());
    assert_eq!(server_bytes, LARGE_TEST_DATA.len());
    assert_eq!(LARGE_TEST_DATA, client_received);
    assert_eq!(LARGE_TEST_DATA, server_received);

    Ok(())
}

#[tokio::test]
async fn send_error() -> Result<(), Box<dyn std::error::Error>> {
    let client = TlsConnector::new(common::client_config()?.build()?);
    let server = TlsAcceptor::new(common::server_config()?.build()?);

    let (server_stream, client_stream) = common::get_streams().await?;
    let client_stream = common::TestStream::new(client_stream);
    let overrides = client_stream.overrides();
    let (mut client, _) =
        common::run_negotiate(&client, client_stream, &server, server_stream).await?;

    // Setup write to fail
    overrides.next_write(Some(Box::new(|_, _, _| {
        Ready(Err(io::Error::from(io::ErrorKind::ConnectionReset)))
    })));

    // Verify write fails
    let result = client.write_all(TEST_DATA).await;
    assert!(result.is_err());

    Ok(())
}

#[tokio::test]
async fn recv_error() -> Result<(), Box<dyn std::error::Error>> {
    let client = TlsConnector::new(common::client_config()?.build()?);
    let server = TlsAcceptor::new(common::server_config()?.build()?);

    let (server_stream, client_stream) = common::get_streams().await?;
    let client_stream = common::TestStream::new(client_stream);
    let overrides = client_stream.overrides();
    let (mut client, _) =
        common::run_negotiate(&client, client_stream, &server, server_stream).await?;

    // Setup read to fail
    overrides.next_read(Some(Box::new(|_, _, _| {
        Ready(Err(io::Error::from(io::ErrorKind::ConnectionReset)))
    })));

    // Verify read fails
    let mut received = [0; 1];
    let result = client.read_exact(&mut received).await;
    assert!(result.is_err());

    Ok(())
}