add missing memory allocation checks

augmentation.c

@@ -72,6 +72,10 @@ static void housekeep_augment_cache_locked(void)
 static augmentation_response *augmentation_response_init(void)
 {
     augmentation_response *response = kzalloc(sizeof(augmentation_response), GFP_KERNEL);
+    if (!response)
+    {
+        return ERR_PTR(-ENOMEM);
+    }
 
     kref_init(&response->kref);
 
@@ -97,22 +101,18 @@ static void augmentation_response_release(struct kref *kref)
     augmentation_response_free(response);
 }
 
-void augmentation_response_get(augmentation_response *response)
+static void augmentation_response_get(augmentation_response *response)
 {
     if (!response)
-    {
         return;
-    }
 
     kref_get(&response->kref);
 }
 
 void augmentation_response_put(augmentation_response *response)
 {
     if (!response)
-    {
         return;
-    }
 
     kref_put(&response->kref, augmentation_response_release);
 }
@@ -121,33 +121,47 @@ static char *get_task_context_key(task_context *ctx)
 {
     int keylen = snprintf(NULL, 0, "%s %d %d %u %s", ctx->cgroup_path, ctx->uid.val, ctx->gid.val, ctx->namespace_ids.mnt, ctx->command_path);
     char *key = kzalloc(keylen + 1, GFP_KERNEL);
+    if (!key)
+    {
+        return ERR_PTR(-ENOMEM);
+    }
     snprintf(key, keylen + 1, "%s %d %d %u %s", ctx->cgroup_path, ctx->uid.val, ctx->gid.val, ctx->namespace_ids.mnt, ctx->command_path);
 
     return key;
 }
 
-static void augmentation_response_cache_set_locked(char *key, augmentation_response *response)
+static int augmentation_response_cache_set_locked(char *key, augmentation_response *response)
 {
     if (!key)
     {
-        pr_err("invalid key");
-        return;
+        return -EINVAL;
+    }
+
+    if (!response)
+    {
+        return -EINVAL;
     }
 
     housekeep_augment_cache_locked();
 
     augmentation_response_cache_entry *new_entry = kzalloc(sizeof(augmentation_response_cache_entry), GFP_KERNEL);
     if (!new_entry)
     {
-        pr_err("memory allocation error");
-        return;
+        return -ENOMEM;
     }
 
-    new_entry->key = strdup(key);
+    new_entry->key = kstrdup(key, GFP_KERNEL);
+    if (!new_entry->key)
+    {
+        kfree(new_entry);
+        return -ENOMEM;
+    }
     new_entry->response = response;
 
     pr_debug("add entry # key[%s]", new_entry->key);
     list_add(&new_entry->list, &augmentation_response_cache);
+
+    return 0;
 }
 
 static augmentation_response *augmentation_response_cache_get_locked(char *key)
@@ -174,34 +188,78 @@ static augmentation_response *augmentation_response_cache_get_locked(char *key)
 augmentation_response *augment_workload()
 {
     augmentation_response *response;
+    void *error;
 
-    augmentation_response_cache_lock();
+    task_context *ctx = get_task_context();
+    if (IS_ERR(ctx))
+        return (void *)ctx;
+
+    char *key = get_task_context_key(ctx);
+    if (IS_ERR(key))
+        return (void *)key;
 
-    char *key = get_task_context_key(get_task_context());
+    augmentation_response_cache_lock();
     response = augmentation_response_cache_get_locked(key);
+
     if (response)
     {
         augmentation_response_get(response);
+        augmentation_response_cache_unlock();
         goto ret;
     }
 
+    augmentation_response_cache_unlock();
+
     response = augmentation_response_init();
+    if (IS_ERR(response))
+    {
+        goto ret;
+    }
+
     command_answer *answer = send_augment_command();
+    if (IS_ERR(answer))
+    {
+        error = answer;
+        goto error;
+    }
     if (answer->error)
-        response->error = strdup(answer->error);
+    {
+        response->error = kstrdup(answer->error, GFP_KERNEL);
+        if (!response->error)
+        {
+            error = ERR_PTR(-ENOMEM);
+            goto error;
+        }
+    }
     else
     {
-        response->response = strdup(answer->answer);
-        augmentation_response_cache_set_locked(key, response);
+        response->response = kstrdup(answer->answer, GFP_KERNEL);
+        if (!response->response)
+        {
+            error = ERR_PTR(-ENOMEM);
+            goto error;
+        }
+
+        augmentation_response_cache_lock();
+        int ret = augmentation_response_cache_set_locked(key, response);
+        if (ret < 0)
+        {
+            error = ERR_PTR(ret);
+            augmentation_response_cache_unlock();
+            goto error;
+        }
         augmentation_response_get(response);
+        augmentation_response_cache_unlock();
     }
 
     free_command_answer(answer);
 
 ret:
     kfree(key);
-
-    augmentation_response_cache_unlock();
-
     return response;
+
+error:
+    kfree(key);
+    augmentation_response_free(response);
+    return error;
 }

augmentation.h

@@ -27,8 +27,12 @@ typedef struct
     struct list_head list;
 } augmentation_response_cache_entry;
 
-void augmentation_response_get(augmentation_response *response);
 void augmentation_response_put(augmentation_response *response);
+/*
+ * augment_workload
+ *
+ * returns an augmentation_response struct pointer or ERR_PTR() on error
+ */
 augmentation_response *augment_workload(void);
 
 #endif

buffer.c

@@ -14,19 +14,26 @@
 buffer_t *buffer_new(int capacity)
 {
     buffer_t *buffer = kmalloc(sizeof(buffer_t), GFP_KERNEL);
+    if (!buffer)
+        return ERR_PTR(-ENOMEM);
     buffer->data = kmalloc(capacity, GFP_KERNEL);
+    if (!buffer->data)
+    {
+        buffer_free(buffer);
+        return ERR_PTR(-ENOMEM);
+    }
     buffer->size = 0;
     buffer->capacity = capacity;
     return buffer;
 }
 
 void buffer_free(buffer_t *buffer)
 {
-    if (buffer)
-    {
-        kfree(buffer->data);
-        kfree(buffer);
-    }
+    if (IS_ERR_OR_NULL(buffer))
+        return;
+
+    kfree(buffer->data);
+    kfree(buffer);
 }
 
 char *buffer_grow(buffer_t *buffer, int len)
@@ -43,6 +50,10 @@ char *buffer_grow(buffer_t *buffer, int len)
         }
 
         buffer->data = krealloc(buffer->data, new_capacity, GFP_KERNEL);
+        if (!buffer->data)
+        {
+            return NULL;
+        }
         buffer->capacity = new_capacity;
     }
 

buffer.h

@@ -11,16 +11,26 @@
 #ifndef buffer_h
 #define buffer_h
 
-typedef struct buffer_t {
+typedef struct buffer_t
+{
   char *data;
   int size;
   int capacity;
 } buffer_t;
 
-
+/*
+ * buffer_new
+ *
+ * returns a buffer_t struct pointer or ERR_PTR() on error
+ */
 buffer_t *buffer_new(int capacity);
 void buffer_free(buffer_t *buffer);
-// returns a pointer to the buffer where the caller can write len long data (possibly resizing the buffer)
+/*
+ * buffer_grow
+ *
+ * returns a buffer_t struct pointer where the caller can write len long data
+ * (possibly resizing the buffer) or NULL on realloc error
+ */
 char *buffer_grow(buffer_t *buffer, int len);
 
 #endif

cert_tools.c

@@ -16,6 +16,7 @@
 #include "cert_tools.h"
 #include "string.h"
 #include "rsa_tools.h"
+#include "task_context.h"
 
 // Define the maximum number of elements inside the cache
 #define MAX_CACHE_LENGTH 64
@@ -51,26 +52,32 @@ size_t linkedlist_length(struct list_head *head)
 
 // add_cert_to_cache adds a certificate chain with a given trust anchor to a linked list. The key will identify this entry.
 // the function is thread safe.
-void add_cert_to_cache(char *key, x509_certificate *cert)
+int add_cert_to_cache(char *key, x509_certificate *cert)
 {
     if (!key)
     {
         pr_err("invalid key");
-        return;
+        return -EINVAL;
     }
 
     cert_with_key *new_entry = kzalloc(sizeof(cert_with_key), GFP_KERNEL);
     if (!new_entry)
     {
-        pr_err("memory allocation error");
-        return;
+        return -ENOMEM;
+    }
+    new_entry->key = kstrdup(key, GFP_KERNEL);
+    if (!new_entry->key)
+    {
+        kfree(new_entry);
+        return -ENOMEM;
     }
-    new_entry->key = strdup(key);
     new_entry->cert = cert;
 
     cert_cache_lock();
     list_add(&new_entry->list, &cert_cache);
     cert_cache_unlock();
+
+    return 0;
 }
 
 // remove_unused_expired_certs_from_cache iterates over the whole cache and tries to clean up the unused/expired certificates.
@@ -199,6 +206,8 @@ bool validate_cert(x509_certificate_validity cert_validity)
 x509_certificate *x509_certificate_init(void)
 {
     x509_certificate *cert = kzalloc(sizeof(x509_certificate), GFP_KERNEL);
+    if (!cert)
+        return ERR_PTR(-ENOMEM);
 
     kref_init(&cert->kref);
 
@@ -207,10 +216,8 @@ x509_certificate *x509_certificate_init(void)
 
 static void x509_certificate_free(x509_certificate *cert)
 {
-    if (!cert)
-    {
+    if (IS_ERR_OR_NULL(cert))
         return;
-    }
 
     free_br_x509_certificate(cert->chain, cert->chain_len);
     free_br_x509_trust_anchors(cert->trust_anchors, cert->trust_anchors_len);
@@ -227,18 +234,16 @@ static void x509_certificate_release(struct kref *kref)
 
 void x509_certificate_get(x509_certificate *cert)
 {
-    if (!cert)
-    {
+    if (IS_ERR_OR_NULL(cert))
         return;
-    }
+
     kref_get(&cert->kref);
 }
 
 void x509_certificate_put(x509_certificate *cert)
 {
-    if (!cert)
-    {
+    if (IS_ERR_OR_NULL(cert))
         return;
-    }
+
     kref_put(&cert->kref, x509_certificate_release);
 }

cert_tools.h

@@ -41,11 +41,16 @@ typedef struct
     struct list_head list;
 } cert_with_key;
 
+/*
+ * x509_certificate_init
+ *
+ * returns an x509_certificate struct pointer or ERR_PTR() on error
+ */
 x509_certificate *x509_certificate_init(void);
 void x509_certificate_get(x509_certificate *cert);
 void x509_certificate_put(x509_certificate *cert);
 
-void add_cert_to_cache(char *key, x509_certificate *cert);
+int add_cert_to_cache(char *key, x509_certificate *cert);
 cert_with_key *find_cert_from_cache(char *key);
 void remove_cert_from_cache(cert_with_key *cert);
 void remove_cert_from_cache_locked(cert_with_key *cert);