malloc: add malloc failure wrapper

ABu33 · Jun 18, 2020 · b49c5b8 · b49c5b8
1 parent b029a40
commit b49c5b8
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 4 deletions.
diff --git a/src/crypto/crypto.c b/src/crypto/crypto.c
@@ -67,6 +67,10 @@ char *gcrypt_encrypt_msg(char *msg, size_t msg_length, int *out_length) {
         msg_length += 0x10U - (msg_length & 0xfU);
 
     char *out = malloc(msg_length);
+    if (!out){
+        fprintf(stderr, "gcry_cipher_encrypt error: not enought memory\n");
+        return NULL;
+    } 
     gcry_error_handle = gcry_cipher_encrypt(gcry_cipher_hd, out, msg_length, msg, msg_length);
     if (gcry_error_handle) {
         fprintf(stderr, "gcry_cipher_encrypt failed:  %s/%s\n",
@@ -84,15 +88,24 @@ char *gcrypt_decrypt_msg(char *msg, size_t msg_length) {
         msg_length += 0x10U - (msg_length & 0xfU);
 
     char *out_buffer = malloc(msg_length);
+    if (!out_buffer){
+        fprintf(stderr, "gcry_cipher_decrypt error: not enought memory\n");
+        return NULL;
+    } 
     gcry_error_handle = gcry_cipher_decrypt(gcry_cipher_hd, out_buffer, msg_length, msg, msg_length);
     if (gcry_error_handle) {
-        fprintf(stderr, "gcry_cipher_encrypt failed:  %s/%s\n",
+        fprintf(stderr, "gcry_cipher_decrypt failed:  %s/%s\n",
                 gcry_strsource(gcry_error_handle),
                 gcry_strerror(gcry_error_handle));
         free(out_buffer);
         return NULL;
     }
     char *out = malloc(strlen(out_buffer) + 1);
+    if (!out){
+        free(out_buffer);
+        fprintf(stderr, "gcry_cipher_decrypt error: not enought memory\n");
+        return NULL;
+    }
     strcpy(out, out_buffer);
     free(out_buffer);
     return out;

diff --git a/src/network/networksocket.c b/src/network/networksocket.c
@@ -99,8 +99,17 @@ void *receive_msg_enc(void *args) {
         recv_string[recv_string_len] = '\0';
 
         char *base64_dec_str = malloc(B64_DECODE_LEN(strlen(recv_string)));
+        if (!base64_dec_str){
+            fprintf(stderr, "Received network error: not enought memory\n");
+            return 0;
+        } 
         int base64_dec_length = b64_decode(recv_string, base64_dec_str, B64_DECODE_LEN(strlen(recv_string)));
         char *dec = gcrypt_decrypt_msg(base64_dec_str, base64_dec_length);
+        if (!dec){
+            free(base64_dec_str);
+            fprintf(stderr, "Received network error: not enought memory\n");
+            return 0;
+        }
 
         printf("Received network message: %s\n", dec);
         free(base64_dec_str);
@@ -134,8 +143,19 @@ int send_string_enc(char *msg) {
     int length_enc;
     size_t msglen = strlen(msg);
     char *enc = gcrypt_encrypt_msg(msg, msglen + 1, &length_enc);
+    if (!enc){
+        fprintf(stderr, "sendto() error: not enought memory\n");
+        pthread_mutex_unlock(&send_mutex);
+        exit(EXIT_FAILURE);
+    } 
 
     char *base64_enc_str = malloc(B64_ENCODE_LEN(length_enc));
+    if (!base64_enc_str){
+        free(enc);
+        fprintf(stderr, "sendto() error: not enought memory\n");
+        pthread_mutex_unlock(&send_mutex);
+        exit(EXIT_FAILURE);
+    } 
     size_t base64_enc_length = b64_encode(enc, length_enc, base64_enc_str, B64_ENCODE_LEN(length_enc));
 
     if (sendto(sock,

diff --git a/src/network/tcpsocket.c b/src/network/tcpsocket.c
@@ -80,32 +80,48 @@ static void client_to_server_state(struct ustream *s) {
 }
 
 static void client_read_cb(struct ustream *s, int bytes) {
-    char *str;
+    char *str, *str_tmp;
     int len = 0;
     uint32_t final_len = sizeof(uint32_t);
     str = malloc(final_len);
+    if (!str) {
+        fprintf(stderr,"not enough memory\n");
+        goto memory_full;
+    }
 
     if ((len = ustream_read(s, str, final_len)) < final_len){//ensure recv sizeof(uint32_t).
         fprintf(stderr,"not complete msg, len:%d, expected len:%u\n", len, final_len);
         goto out;
     }
 
     final_len = ntohl(*(uint32_t *)str) - sizeof(uint32_t);//the final_len in headder includes header itself
-    str = realloc(str, final_len);
+    str_tmp = realloc(str, final_len);
+    if (!str_tmp) {
+        fprintf(stderr,"not enough memory\n");
+        goto out;//On failure, realloc returns a null pointer. The original pointer str remains valid 
+                 //and may need to be deallocated with free() or realloc().
+    }
+    str = str_tmp;
     if ((len = ustream_read(s, str, final_len)) < final_len) {//ensure recv final_len bytes.
         fprintf(stderr,"not complete msg, len:%d, expected len:%u\n", len, final_len);
         goto out;
     }
 
     if (network_config.use_symm_enc) {        
-        char *dec = gcrypt_decrypt_msg(str, final_len);//len of str is final_len        
+        char *dec = gcrypt_decrypt_msg(str, final_len);//len of str is final_len     
+        if (!dec) {
+            fprintf(stderr,"not enough memory\n");
+            goto out;
+        }        
         handle_network_msg(dec);
         free(dec);
     } else {
         handle_network_msg(str);//len of str is final_len
     }
 out:
     free(str);
+memory_full:
+    return;
 }
 
 static void server_cb(struct uloop_fd *fd, unsigned int events) {
@@ -227,10 +243,19 @@ void send_tcp(char *msg) {
         int length_enc;
         size_t msglen = strlen(msg)+1;
         char *enc = gcrypt_encrypt_msg(msg, msglen, &length_enc);
+        if (!enc){
+            fprintf(stderr, "Ustream error: not enought memory\n");
+            return;
+        } 
 
         struct network_con_s *con;
         uint32_t final_len = length_enc + sizeof(final_len);
         char *final_str = malloc(final_len);
+        if (!final_str){
+            free(enc);
+            fprintf(stderr, "Ustream error: not enought memory\n");
+            return;
+        }            
         uint32_t *msg_header = (uint32_t *)final_str;
         *msg_header = htonl(final_len);
         memcpy(final_str+sizeof(final_len), enc, length_enc);
@@ -253,6 +278,10 @@ void send_tcp(char *msg) {
         size_t msglen = strlen(msg) + 1;
         uint32_t final_len = msglen + sizeof(final_len);
         char *final_str = malloc(final_len);
+        if (!final_str){
+            fprintf(stderr, "Ustream error: not enought memory\n");
+            return;
+        } 
         uint32_t *msg_header = (uint32_t *)final_str;
         *msg_header = htonl(final_len);
         memcpy(final_str+sizeof(final_len), msg, msglen);