pk_wrap: use proper macros for sign and verify

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
diff --git a/library/pk_wrap.c b/library/pk_wrap.c
index e17b7dd..9a2a7f0 100644
--- a/library/pk_wrap.c
+++ b/library/pk_wrap.c
@@ -620,226 +620,7 @@
     return ((mbedtls_ecp_keypair *) ctx)->grp.pbits;
 }
 
-#if defined(MBEDTLS_PK_CAN_ECDSA_SOME)
-/* Forward declarations */
-static int ecdsa_verify_wrap(void *ctx, mbedtls_md_type_t md_alg,
-                             const unsigned char *hash, size_t hash_len,
-                             const unsigned char *sig, size_t sig_len);
-
-static int ecdsa_sign_wrap(void *ctx, mbedtls_md_type_t md_alg,
-                           const unsigned char *hash, size_t hash_len,
-                           unsigned char *sig, size_t sig_size, size_t *sig_len,
-                           int (*f_rng)(void *, unsigned char *, size_t), void *p_rng);
-
-#if defined(MBEDTLS_ECP_RESTARTABLE)
-/* Forward declarations */
-static int ecdsa_verify_rs_wrap(void *ctx, mbedtls_md_type_t md_alg,
-                                const unsigned char *hash, size_t hash_len,
-                                const unsigned char *sig, size_t sig_len,
-                                void *rs_ctx);
-
-static int ecdsa_sign_rs_wrap(void *ctx, mbedtls_md_type_t md_alg,
-                              const unsigned char *hash, size_t hash_len,
-                              unsigned char *sig, size_t sig_size, size_t *sig_len,
-                              int (*f_rng)(void *, unsigned char *, size_t), void *p_rng,
-                              void *rs_ctx);
-
-/*
- * Restart context for ECDSA operations with ECKEY context
- *
- * We need to store an actual ECDSA context, as we need to pass the same to
- * the underlying ecdsa function, so we can't create it on the fly every time.
- */
-typedef struct {
-    mbedtls_ecdsa_restart_ctx ecdsa_rs;
-    mbedtls_ecdsa_context ecdsa_ctx;
-} eckey_restart_ctx;
-
-static void *eckey_rs_alloc(void)
-{
-    eckey_restart_ctx *rs_ctx;
-
-    void *ctx = mbedtls_calloc(1, sizeof(eckey_restart_ctx));
-
-    if (ctx != NULL) {
-        rs_ctx = ctx;
-        mbedtls_ecdsa_restart_init(&rs_ctx->ecdsa_rs);
-        mbedtls_ecdsa_init(&rs_ctx->ecdsa_ctx);
-    }
-
-    return ctx;
-}
-
-static void eckey_rs_free(void *ctx)
-{
-    eckey_restart_ctx *rs_ctx;
-
-    if (ctx == NULL) {
-        return;
-    }
-
-    rs_ctx = ctx;
-    mbedtls_ecdsa_restart_free(&rs_ctx->ecdsa_rs);
-    mbedtls_ecdsa_free(&rs_ctx->ecdsa_ctx);
-
-    mbedtls_free(ctx);
-}
-
-static int eckey_verify_rs_wrap(void *ctx, mbedtls_md_type_t md_alg,
-                                const unsigned char *hash, size_t hash_len,
-                                const unsigned char *sig, size_t sig_len,
-                                void *rs_ctx)
-{
-    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    eckey_restart_ctx *rs = rs_ctx;
-
-    /* Should never happen */
-    if (rs == NULL) {
-        return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
-    }
-
-    /* set up our own sub-context if needed (that is, on first run) */
-    if (rs->ecdsa_ctx.grp.pbits == 0) {
-        MBEDTLS_MPI_CHK(mbedtls_ecdsa_from_keypair(&rs->ecdsa_ctx, ctx));
-    }
-
-    MBEDTLS_MPI_CHK(ecdsa_verify_rs_wrap(&rs->ecdsa_ctx,
-                                         md_alg, hash, hash_len,
-                                         sig, sig_len, &rs->ecdsa_rs));
-
-cleanup:
-    return ret;
-}
-
-static int eckey_sign_rs_wrap(void *ctx, mbedtls_md_type_t md_alg,
-                              const unsigned char *hash, size_t hash_len,
-                              unsigned char *sig, size_t sig_size, size_t *sig_len,
-                              int (*f_rng)(void *, unsigned char *, size_t), void *p_rng,
-                              void *rs_ctx)
-{
-    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    eckey_restart_ctx *rs = rs_ctx;
-
-    /* Should never happen */
-    if (rs == NULL) {
-        return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
-    }
-
-    /* set up our own sub-context if needed (that is, on first run) */
-    if (rs->ecdsa_ctx.grp.pbits == 0) {
-        MBEDTLS_MPI_CHK(mbedtls_ecdsa_from_keypair(&rs->ecdsa_ctx, ctx));
-    }
-
-    MBEDTLS_MPI_CHK(ecdsa_sign_rs_wrap(&rs->ecdsa_ctx, md_alg,
-                                       hash, hash_len, sig, sig_size, sig_len,
-                                       f_rng, p_rng, &rs->ecdsa_rs));
-
-cleanup:
-    return ret;
-}
-#endif /* MBEDTLS_ECP_RESTARTABLE */
-#endif /* MBEDTLS_PK_CAN_ECDSA_SOME */
-
-static int eckey_check_pair(const void *pub, const void *prv,
-                            int (*f_rng)(void *, unsigned char *, size_t),
-                            void *p_rng)
-{
-    return mbedtls_ecp_check_pub_priv((const mbedtls_ecp_keypair *) pub,
-                                      (const mbedtls_ecp_keypair *) prv,
-                                      f_rng, p_rng);
-}
-
-static void *eckey_alloc_wrap(void)
-{
-    void *ctx = mbedtls_calloc(1, sizeof(mbedtls_ecp_keypair));
-
-    if (ctx != NULL) {
-        mbedtls_ecp_keypair_init(ctx);
-    }
-
-    return ctx;
-}
-
-static void eckey_free_wrap(void *ctx)
-{
-    mbedtls_ecp_keypair_free((mbedtls_ecp_keypair *) ctx);
-    mbedtls_free(ctx);
-}
-
-static void eckey_debug(const void *ctx, mbedtls_pk_debug_item *items)
-{
-    items->type = MBEDTLS_PK_DEBUG_ECP;
-    items->name = "eckey.Q";
-    items->value = &(((mbedtls_ecp_keypair *) ctx)->Q);
-}
-
-const mbedtls_pk_info_t mbedtls_eckey_info = {
-    MBEDTLS_PK_ECKEY,
-    "EC",
-    eckey_get_bitlen,
-    eckey_can_do,
-#if defined(MBEDTLS_PK_CAN_ECDSA_SOME)
-    ecdsa_verify_wrap,   /* Compatible key structures */
-    ecdsa_sign_wrap,   /* Compatible key structures */
-#if defined(MBEDTLS_ECP_RESTARTABLE)
-    eckey_verify_rs_wrap,
-    eckey_sign_rs_wrap,
-#endif
-#else /* MBEDTLS_PK_CAN_ECDSA_SOME */
-    NULL,
-    NULL,
-#endif /* MBEDTLS_PK_CAN_ECDSA_SOME */
-    NULL,
-    NULL,
-    eckey_check_pair,
-    eckey_alloc_wrap,
-    eckey_free_wrap,
-#if defined(MBEDTLS_ECDSA_C) && defined(MBEDTLS_ECP_RESTARTABLE)
-    eckey_rs_alloc,
-    eckey_rs_free,
-#endif
-    eckey_debug,
-};
-
-/*
- * EC key restricted to ECDH
- */
-static int eckeydh_can_do(mbedtls_pk_type_t type)
-{
-    return type == MBEDTLS_PK_ECKEY ||
-           type == MBEDTLS_PK_ECKEY_DH;
-}
-
-const mbedtls_pk_info_t mbedtls_eckeydh_info = {
-    MBEDTLS_PK_ECKEY_DH,
-    "EC_DH",
-    eckey_get_bitlen,         /* Same underlying key structure */
-    eckeydh_can_do,
-    NULL,
-    NULL,
-#if defined(MBEDTLS_ECDSA_C) && defined(MBEDTLS_ECP_RESTARTABLE)
-    NULL,
-    NULL,
-#endif
-    NULL,
-    NULL,
-    eckey_check_pair,
-    eckey_alloc_wrap,       /* Same underlying key structure */
-    eckey_free_wrap,        /* Same underlying key structure */
-#if defined(MBEDTLS_ECDSA_C) && defined(MBEDTLS_ECP_RESTARTABLE)
-    NULL,
-    NULL,
-#endif
-    eckey_debug,            /* Same underlying key structure */
-};
-#endif /* MBEDTLS_ECP_C */
-
-#if defined(MBEDTLS_PK_CAN_ECDSA_SOME)
-static int ecdsa_can_do(mbedtls_pk_type_t type)
-{
-    return type == MBEDTLS_PK_ECDSA;
-}
-
+#if defined(MBEDTLS_PK_CAN_ECDSA_VERIFY)
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
 /*
  * An ASN.1 encoded signature is a sequence of two ASN.1 integers. Parse one of
@@ -1001,7 +782,9 @@
     return ret;
 }
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
+#endif /* MBEDTLS_PK_CAN_ECDSA_VERIFY */
 
+#if defined(MBEDTLS_PK_CAN_ECDSA_SIGN)
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
 /*
  * Simultaneously convert and move raw MPI from the beginning of a buffer
@@ -1212,7 +995,7 @@
 
     return ret;
 }
-#else
+#else /* MBEDTLS_USE_PSA_CRYPTO */
 static int ecdsa_sign_wrap(void *ctx, mbedtls_md_type_t md_alg,
                            const unsigned char *hash, size_t hash_len,
                            unsigned char *sig, size_t sig_size, size_t *sig_len,
@@ -1223,8 +1006,223 @@
                                          sig, sig_size, sig_len,
                                          f_rng, p_rng);
 }
-#endif
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+#endif /* MBEDTLS_PK_CAN_ECDSA_SIGN */
 
+#if defined(MBEDTLS_PK_CAN_ECDSA_SOME)
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+/* Forward declarations */
+static int ecdsa_verify_rs_wrap(void *ctx, mbedtls_md_type_t md_alg,
+                                const unsigned char *hash, size_t hash_len,
+                                const unsigned char *sig, size_t sig_len,
+                                void *rs_ctx);
+
+static int ecdsa_sign_rs_wrap(void *ctx, mbedtls_md_type_t md_alg,
+                              const unsigned char *hash, size_t hash_len,
+                              unsigned char *sig, size_t sig_size, size_t *sig_len,
+                              int (*f_rng)(void *, unsigned char *, size_t), void *p_rng,
+                              void *rs_ctx);
+
+/*
+ * Restart context for ECDSA operations with ECKEY context
+ *
+ * We need to store an actual ECDSA context, as we need to pass the same to
+ * the underlying ecdsa function, so we can't create it on the fly every time.
+ */
+typedef struct {
+    mbedtls_ecdsa_restart_ctx ecdsa_rs;
+    mbedtls_ecdsa_context ecdsa_ctx;
+} eckey_restart_ctx;
+
+static void *eckey_rs_alloc(void)
+{
+    eckey_restart_ctx *rs_ctx;
+
+    void *ctx = mbedtls_calloc(1, sizeof(eckey_restart_ctx));
+
+    if (ctx != NULL) {
+        rs_ctx = ctx;
+        mbedtls_ecdsa_restart_init(&rs_ctx->ecdsa_rs);
+        mbedtls_ecdsa_init(&rs_ctx->ecdsa_ctx);
+    }
+
+    return ctx;
+}
+
+static void eckey_rs_free(void *ctx)
+{
+    eckey_restart_ctx *rs_ctx;
+
+    if (ctx == NULL) {
+        return;
+    }
+
+    rs_ctx = ctx;
+    mbedtls_ecdsa_restart_free(&rs_ctx->ecdsa_rs);
+    mbedtls_ecdsa_free(&rs_ctx->ecdsa_ctx);
+
+    mbedtls_free(ctx);
+}
+
+static int eckey_verify_rs_wrap(void *ctx, mbedtls_md_type_t md_alg,
+                                const unsigned char *hash, size_t hash_len,
+                                const unsigned char *sig, size_t sig_len,
+                                void *rs_ctx)
+{
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+    eckey_restart_ctx *rs = rs_ctx;
+
+    /* Should never happen */
+    if (rs == NULL) {
+        return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+    }
+
+    /* set up our own sub-context if needed (that is, on first run) */
+    if (rs->ecdsa_ctx.grp.pbits == 0) {
+        MBEDTLS_MPI_CHK(mbedtls_ecdsa_from_keypair(&rs->ecdsa_ctx, ctx));
+    }
+
+    MBEDTLS_MPI_CHK(ecdsa_verify_rs_wrap(&rs->ecdsa_ctx,
+                                         md_alg, hash, hash_len,
+                                         sig, sig_len, &rs->ecdsa_rs));
+
+cleanup:
+    return ret;
+}
+
+static int eckey_sign_rs_wrap(void *ctx, mbedtls_md_type_t md_alg,
+                              const unsigned char *hash, size_t hash_len,
+                              unsigned char *sig, size_t sig_size, size_t *sig_len,
+                              int (*f_rng)(void *, unsigned char *, size_t), void *p_rng,
+                              void *rs_ctx)
+{
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+    eckey_restart_ctx *rs = rs_ctx;
+
+    /* Should never happen */
+    if (rs == NULL) {
+        return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+    }
+
+    /* set up our own sub-context if needed (that is, on first run) */
+    if (rs->ecdsa_ctx.grp.pbits == 0) {
+        MBEDTLS_MPI_CHK(mbedtls_ecdsa_from_keypair(&rs->ecdsa_ctx, ctx));
+    }
+
+    MBEDTLS_MPI_CHK(ecdsa_sign_rs_wrap(&rs->ecdsa_ctx, md_alg,
+                                       hash, hash_len, sig, sig_size, sig_len,
+                                       f_rng, p_rng, &rs->ecdsa_rs));
+
+cleanup:
+    return ret;
+}
+#endif /* MBEDTLS_ECP_RESTARTABLE */
+#endif /* MBEDTLS_PK_CAN_ECDSA_SOME */
+
+static int eckey_check_pair(const void *pub, const void *prv,
+                            int (*f_rng)(void *, unsigned char *, size_t),
+                            void *p_rng)
+{
+    return mbedtls_ecp_check_pub_priv((const mbedtls_ecp_keypair *) pub,
+                                      (const mbedtls_ecp_keypair *) prv,
+                                      f_rng, p_rng);
+}
+
+static void *eckey_alloc_wrap(void)
+{
+    void *ctx = mbedtls_calloc(1, sizeof(mbedtls_ecp_keypair));
+
+    if (ctx != NULL) {
+        mbedtls_ecp_keypair_init(ctx);
+    }
+
+    return ctx;
+}
+
+static void eckey_free_wrap(void *ctx)
+{
+    mbedtls_ecp_keypair_free((mbedtls_ecp_keypair *) ctx);
+    mbedtls_free(ctx);
+}
+
+static void eckey_debug(const void *ctx, mbedtls_pk_debug_item *items)
+{
+    items->type = MBEDTLS_PK_DEBUG_ECP;
+    items->name = "eckey.Q";
+    items->value = &(((mbedtls_ecp_keypair *) ctx)->Q);
+}
+
+const mbedtls_pk_info_t mbedtls_eckey_info = {
+    MBEDTLS_PK_ECKEY,
+    "EC",
+    eckey_get_bitlen,
+    eckey_can_do,
+#if defined(MBEDTLS_PK_CAN_ECDSA_VERIFY)
+    ecdsa_verify_wrap,   /* Compatible key structures */
+#else
+    NULL,
+#endif
+#if defined(MBEDTLS_PK_CAN_ECDSA_SIGN)
+    ecdsa_sign_wrap,   /* Compatible key structures */
+#else
+    NULL,
+#endif
+#if defined(MBEDTLS_PK_CAN_ECDSA_SOME) && defined(MBEDTLS_ECP_RESTARTABLE)
+    eckey_verify_rs_wrap,
+    eckey_sign_rs_wrap,
+#endif
+    NULL,
+    NULL,
+    eckey_check_pair,
+    eckey_alloc_wrap,
+    eckey_free_wrap,
+#if defined(MBEDTLS_ECDSA_C) && defined(MBEDTLS_ECP_RESTARTABLE)
+    eckey_rs_alloc,
+    eckey_rs_free,
+#endif
+    eckey_debug,
+};
+
+/*
+ * EC key restricted to ECDH
+ */
+static int eckeydh_can_do(mbedtls_pk_type_t type)
+{
+    return type == MBEDTLS_PK_ECKEY ||
+           type == MBEDTLS_PK_ECKEY_DH;
+}
+
+const mbedtls_pk_info_t mbedtls_eckeydh_info = {
+    MBEDTLS_PK_ECKEY_DH,
+    "EC_DH",
+    eckey_get_bitlen,         /* Same underlying key structure */
+    eckeydh_can_do,
+    NULL,
+    NULL,
+#if defined(MBEDTLS_ECDSA_C) && defined(MBEDTLS_ECP_RESTARTABLE)
+    NULL,
+    NULL,
+#endif
+    NULL,
+    NULL,
+    eckey_check_pair,
+    eckey_alloc_wrap,       /* Same underlying key structure */
+    eckey_free_wrap,        /* Same underlying key structure */
+#if defined(MBEDTLS_ECDSA_C) && defined(MBEDTLS_ECP_RESTARTABLE)
+    NULL,
+    NULL,
+#endif
+    eckey_debug,            /* Same underlying key structure */
+};
+#endif /* MBEDTLS_ECP_C */
+
+#if defined(MBEDTLS_PK_CAN_ECDSA_SOME)
+static int ecdsa_can_do(mbedtls_pk_type_t type)
+{
+    return type == MBEDTLS_PK_ECDSA;
+}
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
 #if defined(MBEDTLS_ECP_RESTARTABLE)
 static int ecdsa_verify_rs_wrap(void *ctx, mbedtls_md_type_t md_alg,
                                 const unsigned char *hash, size_t hash_len,
@@ -1258,9 +1256,7 @@
         (mbedtls_ecdsa_restart_ctx *) rs_ctx);
 
 }
-#endif /* MBEDTLS_ECP_RESTARTABLE */
 
-#if defined(MBEDTLS_ECP_RESTARTABLE)
 static void *ecdsa_rs_alloc(void)
 {
     void *ctx = mbedtls_calloc(1, sizeof(mbedtls_ecdsa_restart_ctx));
@@ -1278,15 +1274,24 @@
     mbedtls_free(ctx);
 }
 #endif /* MBEDTLS_ECP_RESTARTABLE */
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 const mbedtls_pk_info_t mbedtls_ecdsa_info = {
     MBEDTLS_PK_ECDSA,
     "ECDSA",
     eckey_get_bitlen,     /* Compatible key structures */
     ecdsa_can_do,
-    ecdsa_verify_wrap,
-    ecdsa_sign_wrap,
-#if defined(MBEDTLS_ECP_RESTARTABLE)
+#if defined(MBEDTLS_PK_CAN_ECDSA_VERIFY)
+    ecdsa_verify_wrap,   /* Compatible key structures */
+#else
+    NULL,
+#endif
+#if defined(MBEDTLS_PK_CAN_ECDSA_SIGN)
+    ecdsa_sign_wrap,   /* Compatible key structures */
+#else
+    NULL,
+#endif
+#if defined(MBEDTLS_PK_CAN_ECDSA_SOME) && defined(MBEDTLS_ECP_RESTARTABLE)
     ecdsa_verify_rs_wrap,
     ecdsa_sign_rs_wrap,
 #endif
diff --git a/library/pk_wrap.h b/library/pk_wrap.h
index 84e3009..abfac7a 100644
--- a/library/pk_wrap.h
+++ b/library/pk_wrap.h
@@ -31,6 +31,10 @@
 #include "psa/crypto.h"
 #endif /* MBEDTLS_PSA_CRYPTO_C */
 
+#if defined(MBEDTLS_PK_CAN_ECDSA_VERIFY) || defined(MBEDTLS_PK_CAN_ECDSA_SIGN)
+#define MBEDTLS_PK_CAN_ECDSA_SOME
+#endif
+
 struct mbedtls_pk_info_t {
     /** Public key type */
     mbedtls_pk_type_t type;
@@ -56,11 +60,7 @@
                      int (*f_rng)(void *, unsigned char *, size_t),
                      void *p_rng);
 
-#if defined(MBEDTLS_PK_CAN_ECDSA_VERIFY) || defined(MBEDTLS_PK_CAN_ECDSA_SIGN)
-#define MBEDTLS_PK_CAN_ECDSA_SOME
-#endif
-
-#if defined(MBEDTLS_ECDSA_C) && defined(MBEDTLS_ECP_RESTARTABLE)
+#if defined(MBEDTLS_PK_CAN_ECDSA_SOME) && defined(MBEDTLS_ECP_RESTARTABLE)
     /** Verify signature (restartable) */
     int (*verify_rs_func)(void *ctx, mbedtls_md_type_t md_alg,
                           const unsigned char *hash, size_t hash_len,