malloc: add malloc failure wrapper

This commit is contained in:
twy_2000 2020-06-13 00:01:05 +08:00 committed by Polynomialdivision
parent b029a40b11
commit b49c5b8552
3 changed files with 66 additions and 4 deletions

View file

@ -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;

View file

@ -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,

View file

@ -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);