Make sub-context statically allocated

This makes a mbedtls_pk_context memory-wise equivalent to a
mbedtls_uecc_keypair and removes a dynamic allocation, making the PK layer
zero-cost in terms of memory when PK_SINGLE_TYPE is enabled.

include/mbedtls/pk.h

@@ -49,6 +49,10 @@
 #include "tinycrypt/ecc.h"
 #endif
 
+#if defined(MBEDTLS_PK_SINGLE_TYPE)
+#include "pk_internal.h"
+#endif
+
 #if ( defined(__ARMCC_VERSION) || defined(_MSC_VER) ) && \
     !defined(inline) && !defined(__cplusplus)
 #define inline __inline
@@ -137,17 +141,6 @@ typedef const mbedtls_pk_info_t *mbedtls_pk_handle_t;
 #define MBEDTLS_PK_INVALID_HANDLE ( (mbedtls_pk_handle_t) NULL )
 #endif /* MBEDTLS_PK_SINGLE_TYPE */
 
-/**
- * \brief           Public key container
- */
-typedef struct mbedtls_pk_context
-{
-#if !defined(MBEDTLS_PK_SINGLE_TYPE)
-    mbedtls_pk_handle_t         pk_info; /**< Public key information         */
-#endif
-    void *                      pk_ctx;  /**< Underlying public key context  */
-} mbedtls_pk_context;
-
 #if defined(MBEDTLS_USE_TINYCRYPT)
 typedef struct
 {
@@ -156,6 +149,22 @@ typedef struct
 } mbedtls_uecc_keypair;
 #endif
 
+/**
+ * \brief           Public key container
+ */
+typedef struct mbedtls_pk_context
+{
+#if defined(MBEDTLS_PK_SINGLE_TYPE)
+    /* This is an array to make access to it more uniform with the case where
+     * it's a pointer to void - either way it needs casting before use. */
+    unsigned char pk_ctx[sizeof(
+                         MBEDTLS_PK_INFO_CONTEXT( MBEDTLS_PK_SINGLE_TYPE ) )];
+#else
+    mbedtls_pk_handle_t         pk_info; /**< Public key information         */
+    void *                      pk_ctx;  /**< Underlying public key context  */
+#endif
+} mbedtls_pk_context;
+
 #if defined(MBEDTLS_ECDSA_C) && defined(MBEDTLS_ECP_RESTARTABLE)
 /**
  * \brief           Context for resuming operations
@@ -184,11 +193,18 @@ static inline mbedtls_rsa_context *mbedtls_pk_rsa( const mbedtls_pk_context pk )
 #endif /* MBEDTLS_RSA_C */
 
 #if defined(MBEDTLS_USE_TINYCRYPT)
+#if !defined(MBEDTLS_PK_SINGLE_TYPE)
 static inline mbedtls_uecc_keypair *mbedtls_pk_uecc( const mbedtls_pk_context pk )
 {
     return( (mbedtls_uecc_keypair *) (pk).pk_ctx );
 }
+#else
+/* Go with a macro in order to avoid making a copy of the struct (the argument
+ * is not a pointer so it's passed by value) and then returning an address
+ * inside that copy, which would be undefined behaviour. */
+#define mbedtls_pk_uecc( pk )   ( (mbedtls_uecc_keypair *) (pk).pk_ctx )
 #endif
+#endif /* MBEDTLS_USE_TINYCRYPT */
 
 #if defined(MBEDTLS_ECP_C)
 /**

include/mbedtls/pk_internal.h

@@ -50,6 +50,7 @@
 /* Dummy definition to keep check-names.sh happy - don't uncomment */
 //#define MBEDTLS_PK_INFO_ECKEY
 
+#define MBEDTLS_PK_INFO_ECKEY_CONTEXT           mbedtls_uecc_keypair
 #define MBEDTLS_PK_INFO_ECKEY_TYPE              MBEDTLS_PK_ECKEY
 #define MBEDTLS_PK_INFO_ECKEY_NAME              "EC"
 #define MBEDTLS_PK_INFO_ECKEY_GET_BITLEN        uecc_eckey_get_bitlen
@@ -70,6 +71,7 @@
 /*
  * Helper macros to extract fields from PK types
  */
+#define MBEDTLS_PK_INFO_CONTEXT_T( PK )         PK ## _CONTEXT
 #define MBEDTLS_PK_INFO_TYPE_T( PK )            PK ## _TYPE
 #define MBEDTLS_PK_INFO_NAME_T( PK )            PK ## _NAME
 #define MBEDTLS_PK_INFO_GET_BITLEN_T( PK )      PK ## _GET_BITLEN
@@ -94,6 +96,7 @@
  * field name. This allows to call these macros as
  *    MBEDTLS_PK_INFO_{FIELD}( MBEDTLS_PK_SINGLE_TYPE ).
  * where MBEDTLS_PK_SINGLE_TYPE expands to MBEDTLS_PK_INFO_{TYPE}. */
+#define MBEDTLS_PK_INFO_CONTEXT( PK )         MBEDTLS_PK_INFO_CONTEXT_T( PK )
 #define MBEDTLS_PK_INFO_TYPE( PK )            MBEDTLS_PK_INFO_TYPE_T( PK )
 #define MBEDTLS_PK_INFO_NAME( PK )            MBEDTLS_PK_INFO_NAME_T( PK )
 #define MBEDTLS_PK_INFO_GET_BITLEN( PK )      MBEDTLS_PK_INFO_GET_BITLEN_T( PK )
@@ -113,6 +116,7 @@
 #define MBEDTLS_PK_INFO_DEBUG_FUNC( PK )      MBEDTLS_PK_INFO_DEBUG_FUNC_T( PK )
 #define MBEDTLS_PK_INFO_DEBUG_OMIT( PK )      MBEDTLS_PK_INFO_DEBUG_OMIT_T( PK )
 
+#if !defined(MBEDTLS_PK_SINGLE_TYPE)
 struct mbedtls_pk_info_t
 {
     /** Public key type */
@@ -230,6 +234,7 @@ struct mbedtls_pk_info_t
     MBEDTLS_PK_INFO_DEBUG_FUNC( PK ),        \
 }
 #endif /* MBEDTLS_ECDSA_C && MBEDTLS_ECP_RESTARTABLE */
+#endif /* MBEDTLS_PK_SINGLE_TYPE */
 
 /*
  * Macros to access pk_info

library/pk.c

@@ -731,6 +731,7 @@ static int uecc_eckey_sign_wrap( void *ctx, mbedtls_md_type_t md_alg,
     #undef MAX_SECP256R1_ECDSA_SIG_LEN
 }
 
+#if !defined(MBEDTLS_PK_SINGLE_TYPE)
 static void *uecc_eckey_alloc_wrap( void )
 {
     return( mbedtls_calloc( 1, sizeof( mbedtls_uecc_keypair ) ) );
@@ -744,6 +745,7 @@ static void uecc_eckey_free_wrap( void *ctx )
     mbedtls_platform_zeroize( ctx, sizeof( mbedtls_uecc_keypair ) );
     mbedtls_free( ctx );
 }
+#endif /* MBEDTLS_PK_SINGLE_TYPE */
 
 #if !defined(MBEDTLS_PK_SINGLE_TYPE)
 const mbedtls_pk_info_t mbedtls_uecc_eckey_info =
@@ -1151,20 +1153,6 @@ MBEDTLS_ALWAYS_INLINE static inline int pk_info_check_pair_func(
 #endif
 }
 
-MBEDTLS_ALWAYS_INLINE static inline void *pk_info_ctx_alloc_func(
-    mbedtls_pk_handle_t info )
-{
-    (void) info;
-    return( MBEDTLS_PK_INFO_CTX_ALLOC_FUNC( MBEDTLS_PK_SINGLE_TYPE )( ) );
-}
-
-MBEDTLS_ALWAYS_INLINE static inline void pk_info_ctx_free_func(
-    mbedtls_pk_handle_t info, void *ctx )
-{
-    (void) info;
-    MBEDTLS_PK_INFO_CTX_FREE_FUNC( MBEDTLS_PK_SINGLE_TYPE )( ctx );
-}
-
 MBEDTLS_ALWAYS_INLINE static inline int pk_info_debug_func(
     mbedtls_pk_handle_t info,
     const void *ctx, mbedtls_pk_debug_item *items )
@@ -1301,8 +1289,10 @@ void mbedtls_pk_init( mbedtls_pk_context *ctx )
 
 #if !defined(MBEDTLS_PK_SINGLE_TYPE)
     ctx->pk_info = MBEDTLS_PK_INVALID_HANDLE;
-#endif
     ctx->pk_ctx = NULL;
+#else
+    mbedtls_platform_zeroize( ctx, sizeof( mbedtls_pk_context ) );
+#endif
 }
 
 /*
@@ -1313,8 +1303,10 @@ void mbedtls_pk_free( mbedtls_pk_context *ctx )
     if( ctx == NULL )
         return;
 
+#if !defined(MBEDTLS_PK_SINGLE_TYPE)
     if( MBEDTLS_PK_CTX_IS_VALID( ctx ) )
         pk_info_ctx_free_func( MBEDTLS_PK_CTX_INFO( ctx ), ctx->pk_ctx );
+#endif
 
     mbedtls_platform_zeroize( ctx, sizeof( mbedtls_pk_context ) );
 }
@@ -1404,10 +1396,12 @@ int mbedtls_pk_setup( mbedtls_pk_context *ctx, mbedtls_pk_handle_t info )
         return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
 
     ctx->pk_info = info;
-#endif
 
     if( ( ctx->pk_ctx = pk_info_ctx_alloc_func( info ) ) == NULL )
         return( MBEDTLS_ERR_PK_ALLOC_FAILED );
+#else
+    (void) ctx;
+#endif
 
     return( 0 );
 }