commit 7df2ebbfeaf390cc257c56916c5af94c7785469f
author Damien George <damien.p.george@gmail.com> Tue Oct 31 15:54:15 2017 +1100
committer Damien George <damien.p.george@gmail.com> Wed Dec 13 14:48:53 2017 +1100
tree 1c74ee3ed48e7cbc82d305f1cecd38f9161d3b94
parent 1c52d3e8c64f7d4369bf43fc0eb0cbeb76e32c03

extmod/modussl_mbedtls: Clean up mbedtls state when error during setup.

Without this patch, if the SSL handshake fails (eg the connection was lost)
then the mbedtls state (memory) will never be freed.

extmod/modussl_mbedtls.c

diff --git a/extmod/modussl_mbedtls.c b/extmod/modussl_mbedtls.c
index 69a64d2..197a065 100644
--- a/extmod/modussl_mbedtls.c
+++ b/extmod/modussl_mbedtls.c
@@ -141,8 +141,7 @@
     const byte seed[] = "upy";
     ret = mbedtls_ctr_drbg_seed(&o->ctr_drbg, null_entropy_func/*mbedtls_entropy_func*/, &o->entropy, seed, sizeof(seed));
     if (ret != 0) {
-        printf("ret=%d\n", ret);
-        assert(0);
+        goto cleanup;
     }
 
     ret = mbedtls_ssl_config_defaults(&o->conf,
@@ -150,7 +149,7 @@
                     MBEDTLS_SSL_TRANSPORT_STREAM,
                     MBEDTLS_SSL_PRESET_DEFAULT);
     if (ret != 0) {
-        assert(0);
+        goto cleanup;
     }
 
     mbedtls_ssl_conf_authmode(&o->conf, MBEDTLS_SSL_VERIFY_NONE);
@@ -161,14 +160,14 @@
 
     ret = mbedtls_ssl_setup(&o->ssl, &o->conf);
     if (ret != 0) {
-        assert(0);
+        goto cleanup;
     }
 
     if (args->server_hostname.u_obj != mp_const_none) {
         const char *sni = mp_obj_str_get_str(args->server_hostname.u_obj);
         ret = mbedtls_ssl_set_hostname(&o->ssl, sni);
         if (ret != 0) {
-            assert(0);
+            goto cleanup;
         }
     }
 
@@ -194,13 +193,27 @@
 
     while ((ret = mbedtls_ssl_handshake(&o->ssl)) != 0) {
         if (ret != MBEDTLS_ERR_SSL_WANT_READ && ret != MBEDTLS_ERR_SSL_WANT_WRITE) {
-            //assert(0);
             printf("mbedtls_ssl_handshake error: -%x\n", -ret);
-            mp_raise_OSError(MP_EIO);
+            goto cleanup;
         }
     }
 
     return o;
+
+cleanup:
+    mbedtls_pk_free(&o->pkey);
+    mbedtls_x509_crt_free(&o->cert);
+    mbedtls_x509_crt_free(&o->cacert);
+    mbedtls_ssl_free(&o->ssl);
+    mbedtls_ssl_config_free(&o->conf);
+    mbedtls_ctr_drbg_free(&o->ctr_drbg);
+    mbedtls_entropy_free(&o->entropy);
+
+    if (ret == MBEDTLS_ERR_SSL_ALLOC_FAILED) {
+        mp_raise_OSError(MP_ENOMEM);
+    } else {
+        mp_raise_OSError(MP_EIO);
+    }
 }
 
 STATIC mp_obj_t mod_ssl_getpeercert(mp_obj_t o_in, mp_obj_t binary_form) {