text: /*
   * Copyright (C) 2020 Amazon.com, Inc. or its affiliates.  All Rights Reserved.
   *
   * Permission is hereby granted, free of charge, to any person obtaining a copy of
   * this software and associated documentation files (the "Software"), to deal in
   * the Software without restriction, including without limitation the rights to
   * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
   * the Software, and to permit persons to whom the Software is furnished to do so,
   * subject to the following conditions:
   *
   * The above copyright notice and this permission notice shall be included in all
   * copies or substantial portions of the Software.
   *
   * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
   * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
   * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
   * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
   * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
   * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
   */
  
  /**
   * @file tls_freertos_pkcs11.c
   * @brief TLS transport interface implementations. This implementation uses
   * mbedTLS.
   * @note This file is derived from the tls_freertos.c source file found in the mqtt
   * section of IoT Libraries source code. The file has been modified to support using
   * PKCS #11 when using TLS.
   */
  
  /* Standard includes. */
  #include <string.h>
  
  /* FreeRTOS includes. */
  #include "FreeRTOS.h"
  
  /* FreeRTOS+TCP includes. */
  #include "FreeRTOS_IP.h"
  #include "FreeRTOS_Sockets.h"
  
  /* TLS transport header. */
  #include "tls_freertos_pkcs11.h"
  
  /* FreeRTOS Socket wrapper include. */
  #include "freertos_sockets_wrapper.h"
  
  /* mbedTLS util includes. */
  #include "mbedtls_error.h"
  
  /* PKCS #11 includes. */
  #include "iot_pkcs11_config.h"
  #include "iot_pkcs11.h"
  #include "pkcs11.h"
  #include "iot_pki_utils.h"
  
  /*-----------------------------------------------------------*/
  
  /**
   * @brief Represents string to be logged when mbedTLS returned error
   * does not contain a high-level code.
   */
  static const char * pNoHighLevelMbedTlsCodeStr = "<No-High-Level-Code>";
  
  /**
   * @brief Represents string to be logged when mbedTLS returned error
   * does not contain a low-level code.
   */
  static const char * pNoLowLevelMbedTlsCodeStr = "<No-Low-Level-Code>";
  
  /**
   * @brief Utility for converting the high-level code in an mbedTLS error to string,
   * if the code-contains a high-level code; otherwise, using a default string.
   */
  #define mbedtlsHighLevelCodeOrDefault( mbedTlsCode )        \
      ( mbedtls_strerror_highlevel( mbedTlsCode ) != NULL ) ? \
      mbedtls_strerror_highlevel( mbedTlsCode ) : pNoHighLevelMbedTlsCodeStr
  
  /**
   * @brief Utility for converting the level-level code in an mbedTLS error to string,
   * if the code-contains a level-level code; otherwise, using a default string.
   */
  #define mbedtlsLowLevelCodeOrDefault( mbedTlsCode )        \
      ( mbedtls_strerror_lowlevel( mbedTlsCode ) != NULL ) ? \
      mbedtls_strerror_lowlevel( mbedTlsCode ) : pNoLowLevelMbedTlsCodeStr
  
  /*-----------------------------------------------------------*/
  
  /**
   * @brief Initialize the mbed TLS structures in a network connection.
   *
   * @param[in] pSslContext The SSL context to initialize.
   */
  static void sslContextInit( SSLContext_t * pSslContext );
  
  /**
   * @brief Free the mbed TLS structures in a network connection.
   *
   * @param[in] pSslContext The SSL context to free.
   */
  static void sslContextFree( SSLContext_t * pSslContext );
  
  /**
   * @brief Set up TLS on a TCP connection.
   *
   * @param[in] pNetworkContext Network context.
   * @param[in] pHostName Remote host name, used for server name indication.
   * @param[in] pNetworkCredentials TLS setup parameters.
   *
   * @return #TLS_TRANSPORT_SUCCESS, #TLS_TRANSPORT_INSUFFICIENT_MEMORY, #TLS_TRANSPORT_INVALID_CREDENTIALS,
   * #TLS_TRANSPORT_HANDSHAKE_FAILED, or #TLS_TRANSPORT_INTERNAL_ERROR.
   */
  static TlsTransportStatus_t tlsSetup( NetworkContext_t * pNetworkContext,
                                        const char * pHostName,
                                        const NetworkCredentials_t * pNetworkCredentials );
  
  /**
   * @brief Initialize mbedTLS.
   *
   * @return #TLS_TRANSPORT_SUCCESS, or #TLS_TRANSPORT_INTERNAL_ERROR.
   */
  static TlsTransportStatus_t initMbedtls( void );
  
  /*-----------------------------------------------------------*/
  
  /**
   * @brief Callback that wraps PKCS#11 for pseudo-random number generation.
   *
   * @param[in] pvCtx Caller context.
   * @param[in] pucRandom Byte array to fill with random data.
   * @param[in] xRandomLength Length of byte array.
   *
   * @return Zero on success.
   */
  static int generateRandomBytes( void * pvCtx,
                                  unsigned char * pucRandom,
                                  size_t xRandomLength );
  
  /**
   * @brief Helper for reading the specified certificate object, if present,
   * out of storage, into RAM, and then into an mbedTLS certificate context
   * object.
   *
   * @param[in] pSslContext Caller TLS context.
   * @param[in] pcLabelName PKCS #11 certificate object label.
   * @param[in] xClass PKCS #11 certificate object class.
   * @param[out] pxCertificateContext Certificate context.
   *
   * @return Zero on success.
   */
  static CK_RV readCertificateIntoContext( SSLContext_t * pSslContext,
                                           char * pcLabelName,
                                           CK_OBJECT_CLASS xClass,
                                           mbedtls_x509_crt * pxCertificateContext );
  
  /**
   * @brief Helper for setting up potentially hardware-based cryptographic context.
   *
   * @param Caller context.
   *
   * @return Zero on success.
   */
  static CK_RV initializeClientKeys( SSLContext_t * pxCtx );
  
  /**
   * @brief Sign a cryptographic hash with the private key.
   *
   * @param[in] pvContext Crypto context.
   * @param[in] xMdAlg Unused.
   * @param[in] pucHash Length in bytes of hash to be signed.
   * @param[in] uiHashLen Byte array of hash to be signed.
   * @param[out] pucSig RSA signature bytes.
   * @param[in] pxSigLen Length in bytes of signature buffer.
   * @param[in] piRng Unused.
   * @param[in] pvRng Unused.
   *
   * @return Zero on success.
   */
  static int privateKeySigningCallback( void * pvContext,
                                        mbedtls_md_type_t xMdAlg,
                                        const unsigned char * pucHash,
                                        size_t xHashLen,
                                        unsigned char * pucSig,
                                        size_t * pxSigLen,
                                        int ( * piRng )( void *,
                                                         unsigned char *,
                                                         size_t ),
                                        void * pvRng );
  
  
  /*-----------------------------------------------------------*/
  
  static void sslContextInit( SSLContext_t * pSslContext )
  {
      configASSERT( pSslContext != NULL );
  
      mbedtls_ssl_config_init( &( pSslContext->config ) );
      mbedtls_x509_crt_init( &( pSslContext->rootCa ) );
      mbedtls_x509_crt_init( &( pSslContext->clientCert ) );
      mbedtls_ssl_init( &( pSslContext->context ) );
  
      xInitializePkcs11Session( &( pSslContext->xP11Session ) );
      C_GetFunctionList( &( pSslContext->pxP11FunctionList ) );
  }
  /*-----------------------------------------------------------*/
  
  static void sslContextFree( SSLContext_t * pSslContext )
  {
      configASSERT( pSslContext != NULL );
  
      mbedtls_ssl_free( &( pSslContext->context ) );
      mbedtls_x509_crt_free( &( pSslContext->rootCa ) );
      mbedtls_x509_crt_free( &( pSslContext->clientCert ) );
      mbedtls_ssl_config_free( &( pSslContext->config ) );
  
      pSslContext->pxP11FunctionList->C_CloseSession( pSslContext->xP11Session );
  }
  
  /*-----------------------------------------------------------*/
  
  static TlsTransportStatus_t tlsSetup( NetworkContext_t * pNetworkContext,
                                        const char * pHostName,
                                        const NetworkCredentials_t * pNetworkCredentials )
  {
      TlsTransportStatus_t returnStatus = TLS_TRANSPORT_SUCCESS;
      int mbedtlsError = 0;
      CK_RV xResult = CKR_OK;
  
      configASSERT( pNetworkContext != NULL );
      configASSERT( pHostName != NULL );
      configASSERT( pNetworkCredentials != NULL );
      configASSERT( pNetworkCredentials->pRootCa != NULL );
  
      /* Initialize the mbed TLS context structures. */
      sslContextInit( &( pNetworkContext->sslContext ) );
  
      mbedtlsError = mbedtls_ssl_config_defaults( &( pNetworkContext->sslContext.config ),
                                                  MBEDTLS_SSL_IS_CLIENT,
                                                  MBEDTLS_SSL_TRANSPORT_STREAM,
                                                  MBEDTLS_SSL_PRESET_DEFAULT );
  
      if( mbedtlsError != 0 )
      {
          LogError( ( "Failed to set default SSL configuration: mbedTLSError= %s : %s.",
                      mbedtlsHighLevelCodeOrDefault( mbedtlsError ),
                      mbedtlsLowLevelCodeOrDefault( mbedtlsError ) ) );
  
          /* Per mbed TLS docs, mbedtls_ssl_config_defaults only fails on memory allocation. */
          returnStatus = TLS_TRANSPORT_INSUFFICIENT_MEMORY;
      }
  
      if( returnStatus == TLS_TRANSPORT_SUCCESS )
      {
          /* Set up the certificate security profile, starting from the default value. */
          pNetworkContext->sslContext.certProfile = mbedtls_x509_crt_profile_default;
  
          /* test.mosquitto.org only provides a 1024-bit RSA certificate, which is
           * not acceptable by the default mbed TLS certificate security profile.
           * For the purposes of this demo, allow the use of 1024-bit RSA certificates.
           * This block should be removed otherwise. */
          if( strncmp( pHostName, "test.mosquitto.org", strlen( pHostName ) ) == 0 )
          {
              pNetworkContext->sslContext.certProfile.rsa_min_bitlen = 1024;
          }
  
          /* Set SSL authmode and the RNG context. */
          mbedtls_ssl_conf_authmode( &( pNetworkContext->sslContext.config ),
                                     MBEDTLS_SSL_VERIFY_REQUIRED );
          mbedtls_ssl_conf_rng( &( pNetworkContext->sslContext.config ),
                                generateRandomBytes,
                                &pNetworkContext->sslContext );
          mbedtls_ssl_conf_cert_profile( &( pNetworkContext->sslContext.config ),
                                         &( pNetworkContext->sslContext.certProfile ) );
  
          /* Parse the server root CA certificate into the SSL context. */
          mbedtlsError = mbedtls_x509_crt_parse( &( pNetworkContext->sslContext.rootCa ),
                                                 pNetworkCredentials->pRootCa,
                                                 pNetworkCredentials->rootCaSize );
  
          if( mbedtlsError != 0 )
          {
              LogError( ( "Failed to parse server root CA certificate: mbedTLSError= %s : %s.",
                          mbedtlsHighLevelCodeOrDefault( mbedtlsError ),
                          mbedtlsLowLevelCodeOrDefault( mbedtlsError ) ) );
  
              returnStatus = TLS_TRANSPORT_INVALID_CREDENTIALS;
          }
          else
          {
              mbedtls_ssl_conf_ca_chain( &( pNetworkContext->sslContext.config ),
                                         &( pNetworkContext->sslContext.rootCa ),
                                         NULL );
          }
      }
  
      if( returnStatus == TLS_TRANSPORT_SUCCESS )
      {
          /* Setup the client private key. */
          xResult = initializeClientKeys( &( pNetworkContext->sslContext ) );
  
          if( xResult != CKR_OK )
          {
              LogError( ( "Failed to setup key handling by PKCS #11." ) );
  
              returnStatus = TLS_TRANSPORT_INVALID_CREDENTIALS;
          }
          else
          {
              /* Setup the client certificate. */
              xResult = readCertificateIntoContext( &( pNetworkContext->sslContext ),
                                                    pkcs11configLABEL_DEVICE_CERTIFICATE_FOR_TLS,
                                                    CKO_CERTIFICATE,
                                                    &( pNetworkContext->sslContext.clientCert ) );
  
              if( xResult != CKR_OK )
              {
                  LogError( ( "Failed to get certificate from PKCS #11 module." ) );
  
                  returnStatus = TLS_TRANSPORT_INVALID_CREDENTIALS;
              }
              else
              {
                  mbedtls_ssl_conf_own_cert( &( pNetworkContext->sslContext.config ),
                                             &( pNetworkContext->sslContext.clientCert ),
                                             &( pNetworkContext->sslContext.privKey ) );
              }
          }
      }
  
      if( ( returnStatus == TLS_TRANSPORT_SUCCESS ) && ( pNetworkCredentials->pAlpnProtos != NULL ) )
      {
          /* Include an application protocol list in the TLS ClientHello
           * message. */
          mbedtlsError = mbedtls_ssl_conf_alpn_protocols( &( pNetworkContext->sslContext.config ),
                                                          ( const char ** ) &( pNetworkCredentials->pAlpnProtos ) );
  
          if( mbedtlsError != 0 )
          {
              LogError( ( "Failed to configure ALPN protocol in mbed TLS: mbedTLSError= %s : %s.",
                          mbedtlsHighLevelCodeOrDefault( mbedtlsError ),
                          mbedtlsLowLevelCodeOrDefault( mbedtlsError ) ) );
  
              returnStatus = TLS_TRANSPORT_INTERNAL_ERROR;
          }
      }
  
      if( returnStatus == TLS_TRANSPORT_SUCCESS )
      {
          /* Initialize the mbed TLS secured connection context. */
          mbedtlsError = mbedtls_ssl_setup( &( pNetworkContext->sslContext.context ),
                                            &( pNetworkContext->sslContext.config ) );
  
          if( mbedtlsError != 0 )
          {
              LogError( ( "Failed to set up mbed TLS SSL context: mbedTLSError= %s : %s.",
                          mbedtlsHighLevelCodeOrDefault( mbedtlsError ),
                          mbedtlsLowLevelCodeOrDefault( mbedtlsError ) ) );
  
              returnStatus = TLS_TRANSPORT_INTERNAL_ERROR;
          }
          else
          {
              /* Set the underlying IO for the TLS connection. */
              mbedtls_ssl_set_bio( &( pNetworkContext->sslContext.context ),
                                   pNetworkContext->tcpSocket,
                                   mbedtls_platform_send,
                                   mbedtls_platform_recv,
                                   NULL );
          }
      }
  
      if( returnStatus == TLS_TRANSPORT_SUCCESS )
      {
          /* Enable SNI if requested. */
          if( pNetworkCredentials->disableSni == pdFALSE )
          {
              mbedtlsError = mbedtls_ssl_set_hostname( &( pNetworkContext->sslContext.context ),
                                                       pHostName );
  
              if( mbedtlsError != 0 )
              {
                  LogError( ( "Failed to set server name: mbedTLSError= %s : %s.",
                              mbedtlsHighLevelCodeOrDefault( mbedtlsError ),
                              mbedtlsLowLevelCodeOrDefault( mbedtlsError ) ) );
  
                  returnStatus = TLS_TRANSPORT_INTERNAL_ERROR;
              }
          }
      }
  
      if( returnStatus == TLS_TRANSPORT_SUCCESS )
      {
          /* Perform the TLS handshake. */
          do
          {
              mbedtlsError = mbedtls_ssl_handshake( &( pNetworkContext->sslContext.context ) );
          } while( ( mbedtlsError == MBEDTLS_ERR_SSL_WANT_READ ) ||
                   ( mbedtlsError == MBEDTLS_ERR_SSL_WANT_WRITE ) );
  
          if( mbedtlsError != 0 )
          {
              LogError( ( "Failed to perform TLS handshake: mbedTLSError= %s : %s.",
                          mbedtlsHighLevelCodeOrDefault( mbedtlsError ),
                          mbedtlsLowLevelCodeOrDefault( mbedtlsError ) ) );
  
              returnStatus = TLS_TRANSPORT_HANDSHAKE_FAILED;
          }
      }
  
      if( returnStatus != TLS_TRANSPORT_SUCCESS )
      {
          sslContextFree( &( pNetworkContext->sslContext ) );
      }
      else
      {
          LogInfo( ( "(Network connection %p) TLS handshake successful.",
                     pNetworkContext ) );
      }
  
      return returnStatus;
  }
  
  /*-----------------------------------------------------------*/
  
  static TlsTransportStatus_t initMbedtls( void )
  {
      TlsTransportStatus_t returnStatus = TLS_TRANSPORT_SUCCESS;
  
      /* Set the mutex functions for mbed TLS thread safety. */
      mbedtls_threading_set_alt( mbedtls_platform_mutex_init,
                                 mbedtls_platform_mutex_free,
                                 mbedtls_platform_mutex_lock,
                                 mbedtls_platform_mutex_unlock );
  
      if( returnStatus == TLS_TRANSPORT_SUCCESS )
      {
          LogDebug( ( "Successfully initialized mbedTLS." ) );
      }
  
      return returnStatus;
  }
  
  /*-----------------------------------------------------------*/
  
  static int generateRandomBytes( void * pvCtx,
                                  unsigned char * pucRandom,
                                  size_t xRandomLength )
  {
      /* Must cast from void pointer to conform to mbed TLS API. */
      SSLContext_t * pxCtx = ( SSLContext_t * ) pvCtx;
      CK_RV xResult;
  
      xResult = pxCtx->pxP11FunctionList->C_GenerateRandom( pxCtx->xP11Session, pucRandom, xRandomLength );
  
      if( xResult != CKR_OK )
      {
          LogError( ( "Failed to generate random bytes from the PKCS #11 module." ) );
      }
  
      return xResult;
  }
  
  /*-----------------------------------------------------------*/
  
  static CK_RV readCertificateIntoContext( SSLContext_t * pSslContext,
                                           char * pcLabelName,
                                           CK_OBJECT_CLASS xClass,
                                           mbedtls_x509_crt * pxCertificateContext )
  {
      CK_RV xResult = CKR_OK;
      CK_ATTRIBUTE xTemplate = { 0 };
      CK_OBJECT_HANDLE xCertObj = 0;
  
      /* Get the handle of the certificate. */
      xResult = xFindObjectWithLabelAndClass( pSslContext->xP11Session,
                                              pcLabelName,
                                              xClass,
                                              &xCertObj );
  
      if( ( CKR_OK == xResult ) && ( xCertObj == CK_INVALID_HANDLE ) )
      {
          xResult = CKR_OBJECT_HANDLE_INVALID;
      }
  
      /* Query the certificate size. */
      if( CKR_OK == xResult )
      {
          xTemplate.type = CKA_VALUE;
          xTemplate.ulValueLen = 0;
          xTemplate.pValue = NULL;
          xResult = pSslContext->pxP11FunctionList->C_GetAttributeValue( pSslContext->xP11Session,
                                                                         xCertObj,
                                                                         &xTemplate,
                                                                         1 );
      }
  
      /* Create a buffer for the certificate. */
      if( CKR_OK == xResult )
      {
          xTemplate.pValue = pvPortMalloc( xTemplate.ulValueLen );
  
          if( NULL == xTemplate.pValue )
          {
              xResult = CKR_HOST_MEMORY;
          }
      }
  
      /* Export the certificate. */
      if( CKR_OK == xResult )
      {
          xResult = pSslContext->pxP11FunctionList->C_GetAttributeValue( pSslContext->xP11Session,
                                                                         xCertObj,
                                                                         &xTemplate,
                                                                         1 );
      }
  
      /* Decode the certificate. */
      if( CKR_OK == xResult )
      {
          xResult = mbedtls_x509_crt_parse( pxCertificateContext,
                                            ( const unsigned char * ) xTemplate.pValue,
                                            xTemplate.ulValueLen );
      }
  
      /* Free memory. */
      vPortFree( xTemplate.pValue );
  
      return xResult;
  }
  
  /*-----------------------------------------------------------*/
  
  /**
   * @brief Helper for setting up potentially hardware-based cryptographic context
   * for the client TLS certificate and private key.
   *
   * @param Caller context.
   *
   * @return Zero on success.
   */
  static CK_RV initializeClientKeys( SSLContext_t * pxCtx )
  {
      CK_RV xResult = CKR_OK;
      CK_SLOT_ID * pxSlotIds = NULL;
      CK_ULONG xCount = 0;
      CK_ATTRIBUTE xTemplate[ 2 ];
      mbedtls_pk_type_t xKeyAlgo = ( mbedtls_pk_type_t ) ~0;
  
      /* Get the PKCS #11 module/token slot count. */
      if( CKR_OK == xResult )
      {
          xResult = ( BaseType_t ) pxCtx->pxP11FunctionList->C_GetSlotList( CK_TRUE,
                                                                            NULL,
                                                                            &xCount );
      }
  
      /* Allocate memory to store the token slots. */
      if( CKR_OK == xResult )
      {
          pxSlotIds = ( CK_SLOT_ID * ) pvPortMalloc( sizeof( CK_SLOT_ID ) * xCount );
  
          if( NULL == pxSlotIds )
          {
              xResult = CKR_HOST_MEMORY;
          }
      }
  
      /* Get all of the available private key slot identities. */
      if( CKR_OK == xResult )
      {
          xResult = ( BaseType_t ) pxCtx->pxP11FunctionList->C_GetSlotList( CK_TRUE,
                                                                            pxSlotIds,
                                                                            &xCount );
      }
  
      /* Put the module in authenticated mode. */
      if( CKR_OK == xResult )
      {
          xResult = ( BaseType_t ) pxCtx->pxP11FunctionList->C_Login( pxCtx->xP11Session,
                                                                      CKU_USER,
                                                                      ( CK_UTF8CHAR_PTR ) configPKCS11_DEFAULT_USER_PIN,
                                                                      sizeof( configPKCS11_DEFAULT_USER_PIN ) - 1 );
      }
  
      if( CKR_OK == xResult )
      {
          /* Get the handle of the device private key. */
          xResult = xFindObjectWithLabelAndClass( pxCtx->xP11Session,
                                                  pkcs11configLABEL_DEVICE_PRIVATE_KEY_FOR_TLS,
                                                  CKO_PRIVATE_KEY,
                                                  &pxCtx->xP11PrivateKey );
      }
  
      if( ( CKR_OK == xResult ) && ( pxCtx->xP11PrivateKey == CK_INVALID_HANDLE ) )
      {
          xResult = CK_INVALID_HANDLE;
          LogError( ( "Could not find private key." ) );
      }
  
      /* Query the device private key type. */
      if( xResult == CKR_OK )
      {
          xTemplate[ 0 ].type = CKA_KEY_TYPE;
          xTemplate[ 0 ].pValue = &pxCtx->xKeyType;
          xTemplate[ 0 ].ulValueLen = sizeof( CK_KEY_TYPE );
          xResult = pxCtx->pxP11FunctionList->C_GetAttributeValue( pxCtx->xP11Session,
                                                                   pxCtx->xP11PrivateKey,
                                                                   xTemplate,
                                                                   1 );
      }
  
      /* Map the PKCS #11 key type to an mbedTLS algorithm. */
      if( xResult == CKR_OK )
      {
          switch( pxCtx->xKeyType )
          {
              case CKK_RSA:
                  xKeyAlgo = MBEDTLS_PK_RSA;
                  break;
  
              case CKK_EC:
                  xKeyAlgo = MBEDTLS_PK_ECKEY;
                  break;
  
              default:
                  xResult = CKR_ATTRIBUTE_VALUE_INVALID;
                  break;
          }
      }
  
      /* Map the mbedTLS algorithm to its internal metadata. */
      if( xResult == CKR_OK )
      {
          memcpy( &pxCtx->privKeyInfo, mbedtls_pk_info_from_type( xKeyAlgo ), sizeof( mbedtls_pk_info_t ) );
  
          pxCtx->privKeyInfo.sign_func = privateKeySigningCallback;
          pxCtx->privKey.pk_info = &pxCtx->privKeyInfo;
          pxCtx->privKey.pk_ctx = pxCtx;
      }
  
      /* Free memory. */
      vPortFree( pxSlotIds );
  
      return xResult;
  }
  
  /*-----------------------------------------------------------*/
  
  static int privateKeySigningCallback( void * pvContext,
                                        mbedtls_md_type_t xMdAlg,
                                        const unsigned char * pucHash,
                                        size_t xHashLen,
                                        unsigned char * pucSig,
                                        size_t * pxSigLen,
                                        int ( * piRng )( void *,
                                                         unsigned char *,
                                                         size_t ),
                                        void * pvRng )
  {
      CK_RV xResult = CKR_OK;
      int lFinalResult = 0;
      SSLContext_t * pxTLSContext = ( SSLContext_t * ) pvContext;
      CK_MECHANISM xMech = { 0 };
      CK_BYTE xToBeSigned[ 256 ];
      CK_ULONG xToBeSignedLen = sizeof( xToBeSigned );
  
      /* Unreferenced parameters. */
      ( void ) ( piRng );
      ( void ) ( pvRng );
      ( void ) ( xMdAlg );
  
      /* Sanity check buffer length. */
      if( xHashLen > sizeof( xToBeSigned ) )
      {
          xResult = CKR_ARGUMENTS_BAD;
      }
  
      /* Format the hash data to be signed. */
      if( CKK_RSA == pxTLSContext->xKeyType )
      {
          xMech.mechanism = CKM_RSA_PKCS;
  
          /* mbedTLS expects hashed data without padding, but PKCS #11 C_Sign function performs a hash
           * & sign if hash algorithm is specified.  This helper function applies padding
           * indicating data was hashed with SHA-256 while still allowing pre-hashed data to
           * be provided. */
          xResult = vAppendSHA256AlgorithmIdentifierSequence( ( uint8_t * ) pucHash, xToBeSigned );
          xToBeSignedLen = pkcs11RSA_SIGNATURE_INPUT_LENGTH;
      }
      else if( CKK_EC == pxTLSContext->xKeyType )
      {
          xMech.mechanism = CKM_ECDSA;
          memcpy( xToBeSigned, pucHash, xHashLen );
          xToBeSignedLen = xHashLen;
      }
      else
      {
          xResult = CKR_ARGUMENTS_BAD;
      }
  
      if( CKR_OK == xResult )
      {
          /* Use the PKCS#11 module to sign. */
          xResult = pxTLSContext->pxP11FunctionList->C_SignInit( pxTLSContext->xP11Session,
                                                                 &xMech,
                                                                 pxTLSContext->xP11PrivateKey );
      }
  
      if( CKR_OK == xResult )
      {
          *pxSigLen = sizeof( xToBeSigned );
          xResult = pxTLSContext->pxP11FunctionList->C_Sign( ( CK_SESSION_HANDLE ) pxTLSContext->xP11Session,
                                                             xToBeSigned,
                                                             xToBeSignedLen,
                                                             pucSig,
                                                             ( CK_ULONG_PTR ) pxSigLen );
      }
  
      if( ( xResult == CKR_OK ) && ( CKK_EC == pxTLSContext->xKeyType ) )
      {
          /* PKCS #11 for P256 returns a 64-byte signature with 32 bytes for R and 32 bytes for S.
           * This must be converted to an ASN.1 encoded array. */
          if( *pxSigLen != pkcs11ECDSA_P256_SIGNATURE_LENGTH )
          {
              xResult = CKR_FUNCTION_FAILED;
          }
  
          if( xResult == CKR_OK )
          {
              PKI_pkcs11SignatureTombedTLSSignature( pucSig, pxSigLen );
          }
      }
  
      if( xResult != CKR_OK )
      {
          LogError( ( "Failed to sign message using PKCS #11 with error code %02X.", xResult ) );
      }
  
      return lFinalResult;
  }
  
  /*-----------------------------------------------------------*/
  
  TlsTransportStatus_t TLS_FreeRTOS_Connect( NetworkContext_t * pNetworkContext,
                                             const char * pHostName,
                                             uint16_t port,
                                             const NetworkCredentials_t * pNetworkCredentials,
                                             uint32_t receiveTimeoutMs,
                                             uint32_t sendTimeoutMs )
  {
      TlsTransportStatus_t returnStatus = TLS_TRANSPORT_SUCCESS;
      BaseType_t socketStatus = 0;
  
      if( ( pNetworkContext == NULL ) ||
          ( pHostName == NULL ) ||
          ( pNetworkCredentials == NULL ) )
      {
          LogError( ( "Invalid input parameter(s): Arguments cannot be NULL. pNetworkContext=%p, "
                      "pHostName=%p, pNetworkCredentials=%p.",
                      pNetworkContext,
                      pHostName,
                      pNetworkCredentials ) );
          returnStatus = TLS_TRANSPORT_INVALID_PARAMETER;
      }
      else if( ( pNetworkCredentials->pRootCa == NULL ) )
      {
          LogError( ( "pRootCa cannot be NULL." ) );
          returnStatus = TLS_TRANSPORT_INVALID_PARAMETER;
      }
  
      /* Establish a TCP connection with the server. */
      if( returnStatus == TLS_TRANSPORT_SUCCESS )
      {
          socketStatus = Sockets_Connect( &( pNetworkContext->tcpSocket ),
                                          pHostName,
                                          port,
                                          receiveTimeoutMs,
                                          sendTimeoutMs );
  
          if( socketStatus != 0 )
          {
              LogError( ( "Failed to connect to %s with error %d.",
                          pHostName,
                          socketStatus ) );
              returnStatus = TLS_TRANSPORT_CONNECT_FAILURE;
          }
      }
  
      /* Initialize mbedtls. */
      if( returnStatus == TLS_TRANSPORT_SUCCESS )
      {
          returnStatus = initMbedtls();
      }
  
      /* Perform TLS handshake. */
      if( returnStatus == TLS_TRANSPORT_SUCCESS )
      {
          returnStatus = tlsSetup( pNetworkContext, pHostName, pNetworkCredentials );
      }
  
      /* Clean up on failure. */
      if( returnStatus != TLS_TRANSPORT_SUCCESS )
      {
          if( pNetworkContext->tcpSocket != FREERTOS_INVALID_SOCKET )
          {
              FreeRTOS_closesocket( pNetworkContext->tcpSocket );
          }
      }
      else
      {
          LogInfo( ( "(Network connection %p) Connection to %s established.",
                     pNetworkContext,
                     pHostName ) );
      }
  
      return returnStatus;
  }
  
  /*-----------------------------------------------------------*/
  
  void TLS_FreeRTOS_Disconnect( NetworkContext_t * pNetworkContext )
  {
      BaseType_t tlsStatus = 0;
  
      /* Attempting to terminate TLS connection. */
      tlsStatus = ( BaseType_t ) mbedtls_ssl_close_notify( &( pNetworkContext->sslContext.context ) );
  
      /* Ignore the WANT_READ and WANT_WRITE return values. */
      if( ( tlsStatus != ( BaseType_t ) MBEDTLS_ERR_SSL_WANT_READ ) &&
          ( tlsStatus != ( BaseType_t ) MBEDTLS_ERR_SSL_WANT_WRITE ) )
      {
          if( tlsStatus == 0 )
          {
              LogInfo( ( "(Network connection %p) TLS close-notify sent.",
                         pNetworkContext ) );
          }
          else
          {
              LogError( ( "(Network connection %p) Failed to send TLS close-notify: mbedTLSError= %s : %s.",
                          pNetworkContext,
                          mbedtlsHighLevelCodeOrDefault( tlsStatus ),
                          mbedtlsLowLevelCodeOrDefault( tlsStatus ) ) );
          }
      }
      else
      {
          /* WANT_READ and WANT_WRITE can be ignored. Logging for debugging purposes. */
          LogInfo( ( "(Network connection %p) TLS close-notify sent; ",
                     "received %s as the TLS status can be ignored for close-notify."
                     ( tlsStatus == MBEDTLS_ERR_SSL_WANT_READ ) ? "WANT_READ" : "WANT_WRITE",
                     pNetworkContext ) );
      }
  
      /* Call socket shutdown function to close connection. */
      Sockets_Disconnect( pNetworkContext->tcpSocket );
  
      /* Free mbed TLS contexts. */
      sslContextFree( &( pNetworkContext->sslContext ) );
  
      /* Clear the mutex functions for mbed TLS thread safety. */
      mbedtls_threading_free_alt();
  }
  
  /*-----------------------------------------------------------*/
  
  int32_t TLS_FreeRTOS_recv( NetworkContext_t * pNetworkContext,
                             void * pBuffer,
                             size_t bytesToRecv )
  {
      int32_t tlsStatus = 0;
  
      tlsStatus = ( int32_t ) mbedtls_ssl_read( &( pNetworkContext->sslContext.context ),
                                                pBuffer,
                                                bytesToRecv );
  
      if( ( tlsStatus == MBEDTLS_ERR_SSL_TIMEOUT ) ||
          ( tlsStatus == MBEDTLS_ERR_SSL_WANT_READ ) ||
          ( tlsStatus == MBEDTLS_ERR_SSL_WANT_WRITE ) )
      {
          LogDebug( ( "Failed to read data. However, a read can be retried on this error. "
                      "mbedTLSError= %s : %s.",
                      mbedtlsHighLevelCodeOrDefault( tlsStatus ),
                      mbedtlsLowLevelCodeOrDefault( tlsStatus ) ) );
  
          /* Mark these set of errors as a timeout. The libraries may retry read
           * on these errors. */
          tlsStatus = 0;
      }
      else if( tlsStatus < 0 )
      {
          LogError( ( "Failed to read data: mbedTLSError= %s : %s.",
                      mbedtlsHighLevelCodeOrDefault( tlsStatus ),
                      mbedtlsLowLevelCodeOrDefault( tlsStatus ) ) );
      }
      else
      {
          /* Empty else marker. */
      }
  
      return tlsStatus;
  }
  
  /*-----------------------------------------------------------*/
  
  int32_t TLS_FreeRTOS_send( NetworkContext_t * pNetworkContext,
                             const void * pBuffer,
                             size_t bytesToSend )
  {
      int32_t tlsStatus = 0;
  
      tlsStatus = ( int32_t ) mbedtls_ssl_write( &( pNetworkContext->sslContext.context ),
                                                 pBuffer,
                                                 bytesToSend );
  
      if( ( tlsStatus == MBEDTLS_ERR_SSL_TIMEOUT ) ||
          ( tlsStatus == MBEDTLS_ERR_SSL_WANT_READ ) ||
          ( tlsStatus == MBEDTLS_ERR_SSL_WANT_WRITE ) )
      {
          LogDebug( ( "Failed to send data. However, send can be retried on this error. "
                      "mbedTLSError= %s : %s.",
                      mbedtlsHighLevelCodeOrDefault( tlsStatus ),
                      mbedtlsLowLevelCodeOrDefault( tlsStatus ) ) );
  
          /* Mark these set of errors as a timeout. The libraries may retry send
           * on these errors. */
          tlsStatus = 0;
      }
      else if( tlsStatus < 0 )
      {
          LogError( ( "Failed to send data:  mbedTLSError= %s : %s.",
                      mbedtlsHighLevelCodeOrDefault( tlsStatus ),
                      mbedtlsLowLevelCodeOrDefault( tlsStatus ) ) );
      }
      else
      {
          /* Empty else marker. */
      }
  
      return tlsStatus;
  }
  /*-----------------------------------------------------------*/
  