Free mbedtls structures after use

demo/fw/src/cage.c

@@ -9,7 +9,6 @@
 
 #include "cage.h"
 #include "cage_base64.h"
-#include "logging.h"
 #include "bech32.h"
 #include "rng.h"
 
@@ -19,6 +18,9 @@
 #include "mbedtls/ecp.h"
 #include "mbedtls/ctr_drbg.h"
 
+#include "con_usart.h"
+#define LOG_PRINTF con_printf
+
 
 static enum ca_error parse_stanza(struct ca_keystore *ks, const char *stanza_head, size_t len, unsigned char file_key[16]);
 static enum ca_error parse_stanza_x25519(struct ca_keystore *ks, size_t nargs, const char **args, size_t body_len, const unsigned char *body, unsigned char file_key[16]);
@@ -28,6 +30,10 @@ void ca_keystore_init(struct ca_keystore *ks) {
     mbedtls_ecp_keypair_init(&(ks->x25519_kp));
 }
 
+void ca_keystore_free(struct ca_keystore *ks) {
+    mbedtls_ecp_keypair_free(&(ks->x25519_kp));
+}
+
 enum ca_error ca_keystore_load_x25519_private_key(struct ca_keystore *ks, const unsigned char buf[32]) {
 
     /*
@@ -233,6 +239,7 @@ errout:
 
 enum ca_error parse_stanza_x25519(struct ca_keystore *ks, size_t nargs, const char **args, size_t body_len, const unsigned char *body, unsigned char file_key[16]) {
     int mbedtls_rc = -1;
+    enum ca_error err = CA_ERR_CORRUPTED_STATE;
 
     /* body length: 16 bytes ciphertext of 128 bit file key + 16 bytes Poly1305 tag */
     if (nargs != 1 || body_len != 32) {
@@ -266,13 +273,15 @@ enum ca_error parse_stanza_x25519(struct ca_keystore *ks, size_t nargs, const ch
     mbedtls_rc = mbedtls_ecp_group_load(&x25519_grp, MBEDTLS_ECP_DP_CURVE25519);
     if (mbedtls_rc) {
         LOG_PRINTF("Error: mbedtls_ecp_group_load: rc=-0x%x\n", -mbedtls_rc);
-        return CA_ERR_KEY_WRAPPING_DECRYPTION_FAILURE;
+        err = CA_ERR_KEY_WRAPPING_DECRYPTION_FAILURE;
+        goto err_free_grp;
     }
 
     mbedtls_rc = mbedtls_ecp_point_read_binary(&x25519_grp, &ecp_ciphertext_pt, ecp_ciphertext, sizeof(ecp_ciphertext));
     if (mbedtls_rc) {
         LOG_PRINTF("Error: mbedtls_ecp_point_read_binary: rc=-0x%x\n", -mbedtls_rc);
-        return CA_ERR_KEY_WRAPPING_DECRYPTION_FAILURE;
+        err = CA_ERR_KEY_WRAPPING_DECRYPTION_FAILURE;
+        goto err_free_grp;
     }
 
     /*
@@ -290,7 +299,8 @@ enum ca_error parse_stanza_x25519(struct ca_keystore *ks, size_t nargs, const ch
              mbedtls_ctr_drbg_random, g_drbg_ctx);
     if (mbedtls_rc) {
         LOG_PRINTF("Error: mbedtls_ecp_mul: rc=-0x%x\n", -mbedtls_rc);
-        return CA_ERR_MBEDTLS_ERROR;
+        err = CA_ERR_MBEDTLS_ERROR;
+        goto err_free_dec;
     }
 
 
@@ -310,18 +320,21 @@ enum ca_error parse_stanza_x25519(struct ca_keystore *ks, size_t nargs, const ch
 
     if (mbedtls_rc) {
         LOG_PRINTF("Error: mbedtls_ecp_point_write_binary: rc=-0x%x\n", -mbedtls_rc);
-        return CA_ERR_MBEDTLS_ERROR;
+        err = CA_ERR_MBEDTLS_ERROR;
+        goto err_free_dec;
     }
 
     if (ecp_plaintext_len != 32) {
         LOG_PRINTF("Error: parse_stanza_x25519: (ecp_plaintext_len == %d) != 32 (bug, should not happen!)\n", ecp_plaintext_len);
-        return CA_ERR_CORRUPTED_STATE;
+        err = CA_ERR_CORRUPTED_STATE;
+        goto err_free_dec;
     }
 
     const mbedtls_md_info_t *sha256_md = mbedtls_md_info_from_type(MBEDTLS_MD_SHA256);
     if (!sha256_md) {
         LOG_PRINTF("Error: mbedtls_md_info_from_type: cannot load SHA256\n");
-        return CA_ERR_MBEDTLS_ERROR;
+        err = CA_ERR_MBEDTLS_ERROR;
+        goto err_free_dec;
     }
 
     unsigned char salt[64];
@@ -331,11 +344,13 @@ enum ca_error parse_stanza_x25519(struct ca_keystore *ks, size_t nargs, const ch
                                     MBEDTLS_ECP_PF_UNCOMPRESSED, &pbin_olen, salt+32, 32);
     if (mbedtls_rc) {
         LOG_PRINTF("Error: mbedtls_ecp_point_write_binary: rc=-0x%x\n", -mbedtls_rc);
-        return CA_ERR_MBEDTLS_ERROR;
+        err = CA_ERR_MBEDTLS_ERROR;
+        goto err_free_dec;
     }
     if (pbin_olen != 32) {
         LOG_PRINTF("Error: mbedtls_ecp_point_write_binary returned invalid output len: %d != 32\n", pbin_olen);
-        return CA_ERR_CORRUPTED_STATE;
+        err = CA_ERR_CORRUPTED_STATE;
+        goto err_free_dec;
     }
 
     /*
@@ -370,18 +385,31 @@ enum ca_error parse_stanza_x25519(struct ca_keystore *ks, size_t nargs, const ch
     mbedtls_rc = mbedtls_chachapoly_setkey(&cp_ctx, wrap_key);
     if (mbedtls_rc) {
         LOG_PRINTF("Error: mbedtls_chachapoly_setkey: rc=-0x%x\n", -mbedtls_rc);
-        return CA_ERR_MBEDTLS_ERROR;
+        err = CA_ERR_MBEDTLS_ERROR;
+        goto err_free_cp;
     }
 
     const unsigned char nonce[12] = { 0 };
     mbedtls_rc = mbedtls_chachapoly_auth_decrypt(&cp_ctx, 16, nonce, NULL, 0, (const unsigned char*)(body + body_len - 16), body, file_key);
     if (mbedtls_rc) {
         LOG_PRINTF("Error: mbedtls_chachapoly_auth_decrypt: rc=-0x%x\n", -mbedtls_rc);
-        return CA_ERR_MBEDTLS_ERROR;
+        err = CA_ERR_MBEDTLS_ERROR;
+        goto err_free_cp;
     }
 
+    err = CA_ERR_SUCCESS;
 
-    return CA_ERR_SUCCESS;
+err_free_cp:
+    mbedtls_chachapoly_free(&cp_ctx);
+
+err_free_dec:
+    mbedtls_ecp_point_free(&ecp_dec_res_pt);
+
+err_free_grp:
+    mbedtls_ecp_group_free(&x25519_grp);
+    mbedtls_ecp_point_free(&ecp_ciphertext_pt);
+
+    return err;
 }
 
 enum ca_error check_file_key(const unsigned char *buf, size_t buflen, const unsigned char file_key[16]) {
@@ -441,6 +469,7 @@ enum ca_error check_file_key(const unsigned char *buf, size_t buflen, const unsi
 }
 
 enum ca_error stream_decrypt(unsigned char *out, size_t outlen, size_t *out_written, const unsigned char *in, size_t inlen, const unsigned char file_key[16]) {
+    enum ca_error err = CA_ERR_CORRUPTED_STATE;
 
     const char *found = strstr((const char *)in, "\n--- ");
     if (!found) {
@@ -517,7 +546,8 @@ enum ca_error stream_decrypt(unsigned char *out, size_t outlen, size_t *out_writ
     mbedtls_rc = mbedtls_chachapoly_setkey(&cp_ctx, stream_key);
     if (mbedtls_rc) {
         LOG_PRINTF("Error: mbedtls_chachapoly_setkey: rc=-0x%x\n", -mbedtls_rc);
-        return CA_ERR_MBEDTLS_ERROR;
+        err = CA_ERR_MBEDTLS_ERROR;
+        goto err_free_cp;
     }
 
     in += 16;
@@ -552,14 +582,20 @@ enum ca_error stream_decrypt(unsigned char *out, size_t outlen, size_t *out_writ
                 data_len, (unsigned char *)&nonce, NULL, 0, tag, inp, outp);
         if (mbedtls_rc) {
             LOG_PRINTF("Error: mbedtls_chachapoly_auth_decrypt: i=%zx rc=-0x%x\n", i, -mbedtls_rc);
-            return CA_ERR_MBEDTLS_ERROR;
+            err = CA_ERR_MBEDTLS_ERROR;
+            goto err_free_cp;
         }
 
         *out_written += data_len;
         block_num += 1;
     }
 
-    return CA_ERR_SUCCESS;
+    err = CA_ERR_SUCCESS;
+
+err_free_cp:
+    mbedtls_chachapoly_free(&cp_ctx);
+
+    return err;
 }
 
 #ifdef CAGE_DEBUG

demo/fw/src/cage.h

@@ -33,6 +33,7 @@ struct ca_keystore {
 
 enum ca_error parse_age_buf(struct ca_keystore *ks, const char *buf, size_t buflen, unsigned char file_key[16]);
 void ca_keystore_init(struct ca_keystore *ks);
+void ca_keystore_free(struct ca_keystore *ks);
 enum ca_error ca_keystore_load_x25519_private_key(struct ca_keystore *ks, const unsigned char buf[32]);
 enum ca_error stream_decrypt(unsigned char *out, size_t outlen, size_t *out_written, const unsigned char *in, size_t inlen, const unsigned char file_key[16]);
 

demo/fw/src/main.c

@@ -338,11 +338,12 @@ int main(void)
         spif_printf(&spif, "loop #%u  \033[4%dmloop #%u  \033[10%dmloop #%u\033[0m\n", i, k%8, (k+1)%8);
         */
         i++;
-        k = (k + 1) % 40;
+        k = (k + 1) % 5;
-        if (k == 0) {
-            spif_capture_read(&spif, sizeof(payload_buf), payload_buf);
 
-            con_printf_blocking("\r\n");
+        spif_capture_read(&spif, sizeof(payload_buf), payload_buf);
+
+        con_printf_blocking("\r\n");
+        if (k == 0) {
             con_printf_blocking("\r\n");
             con_printf_blocking("buffer capture:\r\n");
             for (size_t e=0; e<256; e += 27) {
@@ -367,40 +368,46 @@ int main(void)
                 }
             }
             con_printf_blocking("\033[0m\r\n");
+        }
 
-            unpack_px_data(payload_buf, sizeof(payload_buf));
+        unpack_px_data(payload_buf, sizeof(payload_buf));
-
-            size_t payload_len = be32toh(((uint32_t *)payload_buf)[0]);
-            if (payload_len < 0 || payload_len > sizeof(payload_buf) - sizeof(uint32_t)) {
-                con_printf_blocking("Invalid payload size %zx\r\n", payload_len);
-                continue;
-            }
-
-            unsigned char file_key[16];
-            err = parse_age_buf(&ks, payload_buf + sizeof(uint32_t), payload_len+1, file_key);
-            if (err) {
-                con_printf_blocking("Error parsing payload age header: %d\r\n", err);
-                continue;
-            }
-
-            size_t decrypted_size = 0;
-            err = stream_decrypt(dec_buf, sizeof(dec_buf), &decrypted_size, payload_buf + sizeof(uint32_t), payload_len, file_key);
-            if (err) {
-                con_printf_blocking("Error decrypting payload: %d\r\n", err);
-                continue;
-            }
-
-            assert (decrypted_size > 0 && decrypted_size < sizeof(dec_buf));
-            dec_buf[decrypted_size-1] = '\0'; /* overwrite trailing \n */
-            con_printf_blocking("decrypted payload: %s\r\n", dec_buf);
 
+        size_t payload_len = be32toh(((uint32_t *)payload_buf)[0]);
+        if (payload_len < 0 || payload_len > sizeof(payload_buf) - sizeof(uint32_t)) {
+            con_printf_blocking("Invalid payload size %zx\r\n", payload_len);
             spif_printf(&spif, "\033[H\033[0m\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\033[H");
-            spif_printf(&spif, "\033[1;91mDecrypted data:\033[0m\n\n%s\n", dec_buf);
+            spif_printf(&spif, "\033[1;91mDecryption error!\033[0m\n");
+            continue;
         }
 
-        for (int j=0; j<1000000; j++) {
+        unsigned char file_key[16];
-            /* pass */
+        err = parse_age_buf(&ks, payload_buf + sizeof(uint32_t), payload_len+1, file_key);
+        if (err) {
+            con_printf_blocking("Error parsing payload age header: %d\r\n", err);
+            spif_printf(&spif, "\033[H\033[0m\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\033[H");
+            spif_printf(&spif, "\033[1;91mDecryption error!\033[0m\n");
+            continue;
         }
+
+        size_t decrypted_size = 0;
+        err = stream_decrypt(dec_buf, sizeof(dec_buf), &decrypted_size, payload_buf + sizeof(uint32_t), payload_len, file_key);
+        if (err) {
+            con_printf_blocking("Error decrypting payload: %d\r\n", err);
+            spif_printf(&spif, "\033[H\033[0m\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\033[H");
+            spif_printf(&spif, "\033[1;91mDecryption error!\033[0m\n");
+            continue;
+        }
+
+        assert (decrypted_size > 0 && decrypted_size < sizeof(dec_buf));
+        dec_buf[decrypted_size-1] = '\0'; /* overwrite trailing \n */
+        con_printf_blocking("decrypted payload: %s\r\n", dec_buf);
+
+        spif_printf(&spif, "\033[H\033[0m\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\033[H");
+        spif_printf(&spif, "\033[1;92mDecrypted data:\033[0m\n\n%s\n", dec_buf);
+
+        //for (int j=0; j<1000000; j++) {
+            /* pass */
+        //}
     }
 
     return 0;