diff --git a/src/crypto/crypto.c b/src/crypto/crypto.c index 2c34e57..f488c8e 100644 --- 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 index ef9c9c1..cf22f02 100644 --- 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 index 3eb266a..68e1b9f 100644 --- a/src/network/tcpsocket.c +++ b/src/network/tcpsocket.c @@ -80,10 +80,14 @@ 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); @@ -91,14 +95,24 @@ static void client_read_cb(struct ustream *s, int bytes) { } 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 { @@ -106,6 +120,8 @@ static void client_read_cb(struct ustream *s, int bytes) { } 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);