fix: crypto - check for oom in base32 and base64

mlawren · May 29, 2023 · 7f50e91 · 7f50e91
1 parent 49f5c46
commit 7f50e91
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 10 deletions.
diff --git a/src/crypto/base32.c b/src/crypto/base32.c
@@ -14,7 +14,10 @@ static const char base32_chars[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZ234567";
 uint8_t* base32_encode(const uint8_t* src, size_t len, size_t* out_len) {
     *out_len = ((len + 4) / 5) * 8;
     uint8_t* encoded = malloc(*out_len + 1);
-    encoded[*out_len] = '\0';
+    if (encoded == NULL) {
+        *out_len = 0;
+        return NULL;
+    }
 
     for (size_t i = 0, j = 0; i < len;) {
         uint32_t octet0 = i < len ? src[i++] : 0;
@@ -40,6 +43,7 @@ uint8_t* base32_encode(const uint8_t* src, size_t len, size_t* out_len) {
         }
     }
 
+    encoded[*out_len] = '\0';
     return encoded;
 }
 
@@ -49,6 +53,11 @@ uint8_t* base32_decode(const uint8_t* src, size_t len, size_t* out_len) {
     }
     *out_len = len * 5 / 8;
     uint8_t* decoded = malloc(*out_len);
+    if (decoded == NULL) {
+        *out_len = 0;
+        return NULL;
+    }
+
     size_t bits = 0, value = 0, count = 0;
     for (size_t i = 0; i < len; i++) {
         uint8_t c = src[i];

diff --git a/src/crypto/base64.c b/src/crypto/base64.c
@@ -12,12 +12,15 @@ static const char base64_chars[] =
 
 uint8_t* base64_encode(const uint8_t* src, size_t len, size_t* out_len) {
     uint8_t* encoded = NULL;
-    size_t i, j, enc_len;
+    size_t i, j;
     uint32_t octets;
 
-    enc_len = ((len + 2) / 3) * 4;
-    encoded = (uint8_t*)malloc(enc_len + 1);
-    encoded[enc_len] = '\0';
+    *out_len = ((len + 2) / 3) * 4;
+    encoded = malloc(*out_len + 1);
+    if (encoded == NULL) {
+        *out_len = 0;
+        return NULL;
+    }
 
     for (i = 0, j = 0; i < len; i += 3, j += 4) {
         octets =
@@ -29,13 +32,13 @@ uint8_t* base64_encode(const uint8_t* src, size_t len, size_t* out_len) {
     }
 
     if (len % 3 == 1) {
-        encoded[enc_len - 1] = '=';
-        encoded[enc_len - 2] = '=';
+        encoded[*out_len - 1] = '=';
+        encoded[*out_len - 2] = '=';
     } else if (len % 3 == 2) {
-        encoded[enc_len - 1] = '=';
+        encoded[*out_len - 1] = '=';
     }
 
-    *out_len = enc_len;
+    encoded[*out_len] = '\0';
     return encoded;
 }
 
@@ -65,8 +68,9 @@ uint8_t* base64_decode(const uint8_t* src, size_t len, size_t* out_len) {
     }
 
     *out_len = (len / 4) * 3 - padding;
-    uint8_t* decoded = (uint8_t*)malloc(*out_len);
+    uint8_t* decoded = malloc(*out_len);
     if (decoded == NULL) {
+        *out_len = 0;
         return NULL;
     }