Refactored main into different helper functions

client/include/utils.h

@@ -33,7 +33,6 @@ static inline char *get_path_from_url(const char *url) {
   return path ? (char *)path : NULL;
 }
 
-int split_fraction_links(char *data, char *data_arr[], int maxlines);
 void print_hex(const unsigned char *data, size_t size);
 
 void init_random(void);

client/src/fraction.c

@@ -141,6 +141,7 @@ int calc_crc(fraction_t *frac){
 }
 
 void fraction_free(fraction_t *fraction) {
+  if (fraction == NULL) return;
   free(fraction->data);
   fraction->magic = 0;
   fraction->index = 0;

client/src/main.c

@@ -12,14 +12,6 @@
 #define SERVER_IP "127.0.0.1"
 #define SERVER_PORT "8000"
 
-/* Helper functions to assist with cleanup, I hate cleanup */
-static void cleanup_string_array(char **array, int n_elem) {
-  for (int i = 0; i < n_elem; i++) {
-    free(array[i]);
-  }
-  free(array);
-}
-
 static void cleanup_fraction_array(fraction_t *array, int n_elem) {
   for (int i = 0; i < n_elem; i++) {
     fraction_free(&array[i]);
@@ -35,176 +27,186 @@ static int do_connect(void) {
 
   if (h_getaddrinfo(SERVER_IP, SERVER_PORT, &hints, &ainfo) != 0) {
     log_error("Failed to resolve server address");
-    return EXIT_FAILURE;
+    return -1;
   }
 
   printf("Connecting to: %s:%s\n", SERVER_IP, SERVER_PORT);
   sfd = create_sock_and_conn(ainfo);
   if (sfd == -1) {
     log_error("Failed to create socket and connect");
-    return EXIT_FAILURE;
+    return -1;
   }
   freeaddrinfo(ainfo);
 
   return sfd;
 }
 
-int main(void) {
-  int sfd = -1; // to be extra professional
-
-  http_res_t http_fraction_res = {0};
-  http_res_t http_post_res = {0};
-
-  char **fraction_links = NULL;
-  fraction_t *fractions = NULL;
-
-  uint8_t *module = NULL;
-  ssize_t module_size;
-
+static uint8_t *get_aes_key(int sfd, size_t *key_length) {
   EVP_PKEY *pkey = NULL;
-  char *private_key = NULL;
   char *public_key = NULL;
 
+  http_res_t http_post_res = {0};
+
   unsigned char *b64_decoded_aes_key;
   unsigned char *aes_key = NULL;
   size_t key_len = 0;
-
-
-  if (geteuid() != 0) {
-    log_error("This program needs to be run as root!");
-    exit(1);
-  }
-
-  init_random();
-  log_set_level(LOG_DEBUG);
-
-  sfd = do_connect();
 
   pkey = generate_rsa_private_key();
   if (pkey == NULL) {
-    return EXIT_FAILURE;
+    return NULL;
   }
+
   public_key = write_rsa_public_key(pkey);
   if (public_key == NULL) {
-    return EXIT_FAILURE;
+    EVP_PKEY_free(pkey);
+    return NULL;
   }
 
   /* Receive and decrypt AES key from server */
-  if (http_post(sfd, "/", "application/octect-stream", public_key, &http_post_res) !=
-    HTTP_SUCCESS) {
+  if (http_post(sfd, "/", "application/octect-stream", public_key,
+                &http_post_res) != HTTP_SUCCESS) {
     log_error("Failed to send RSA public key");
-    goto cleanup;
+    free(public_key);
+    EVP_PKEY_free(pkey);
+    return NULL;
   }
 
   log_info("Base64 encoded key: %s", http_post_res.data);
   base64_decode(http_post_res.data, &b64_decoded_aes_key, &key_len);
   log_info("Key size (decoded): %zu", key_len);
-
-
+
   aes_key = decrypt_rsa_oaep_evp(pkey, b64_decoded_aes_key, key_len, &key_len);
   if (aes_key == NULL) {
     log_error("Failed to decrypt data from server");
-    goto cleanup;
+    free(b64_decoded_aes_key);
+    http_free(&http_post_res);
+    free(public_key);
+    EVP_PKEY_free(pkey);
+    return NULL;
   }
 
-  print_hex(aes_key, key_len);
+  free(b64_decoded_aes_key);
+  http_free(&http_post_res);
+  free(public_key);
+  EVP_PKEY_free(pkey);
+
+  *key_length = key_len;
+  return aes_key;
+}
+
+static fraction_t *fetch_fractions(int sfd, int *fraction_count) {
+  http_res_t http_fraction_res = {0};
+
+  fraction_t *fractions = NULL;
+
+  int i, num_links;
+  char *line;
 
   if (http_get(sfd, "/", &http_fraction_res) != HTTP_SUCCESS) {
     log_error("Failed to retrieve fraction links");
-    goto cleanup;
   }
 
   log_debug("Retrieved fraction links");
 
-  int num_links = count_lines(http_fraction_res.data) + 1;
-  log_debug("%d links found", num_links);
-
-  fraction_links = calloc(num_links, sizeof(char *));
-  if (!fraction_links) {
-    log_error("Failed to allocate memory for fraction links");
-    goto cleanup;
-  }
+  num_links = count_lines(http_fraction_res.data) + 1;
 
-  int lines_read = split_fraction_links(http_fraction_res.data, fraction_links, num_links);
-  if (lines_read < 0) {
-    log_error("Failed to split fraction links");
-    goto cleanup;
-  }
+  log_debug("%d links found", num_links);
 
-  fractions = malloc(lines_read * sizeof(fraction_t));
+  fractions = calloc(num_links, sizeof(fraction_t));
   if (!fractions) {
     log_error("Failed to allocate memory for fractions");
-    goto cleanup;
+    http_free(&http_fraction_res);
+    return NULL;
   }
 
-  for (int i = 0; i < lines_read; i++) {
-    log_debug("Downloading %s", fraction_links[i]);
-    if (download_fraction(sfd, fraction_links[i], &fractions[i]) != 0) {
+  i = 0;
+  line = strtok(http_fraction_res.data, "\n");
+  while (line != NULL && i < num_links) {
+    log_debug("Downloading %s", line);
+
+    if (download_fraction(sfd, line, &fractions[i]) != 0) {
       log_error("Failed to download fraction");
-      goto cleanup;
+
+      // we have to cleanup only until i because the other fractions have not
+      // been initialized anyhow
+      http_free(&http_fraction_res);
+      cleanup_fraction_array(fractions, i);
+      return NULL;
     }
+
+    i++;
+    line = strtok(NULL, "\n");
   }
 
-  log_info("Downloaded fractions");
+  http_free(&http_fraction_res);
+  *fraction_count = i;
+  return fractions;
+}
+
+int main(void) {
+  int sfd = -1; // to be extra professional
+
+  unsigned char *aes_key = NULL;
+  size_t key_len = 0;
 
-  qsort(fractions, lines_read, sizeof(fraction_t), compare_fractions);
+  fraction_t *fractions;
+  int fraction_count;
 
-  for (int i = 0; i < lines_read; i++) {
-    print_fraction(fractions[i]);
+  uint8_t *module = NULL;
+  ssize_t module_size;
+
+  if (geteuid() != 0) {
+    log_error("This program needs to be run as root!");
+    exit(1);
   }
 
-  module = decrypt_lkm(fractions, num_links, &module_size, aes_key);
+  init_random();
+  log_set_level(LOG_DEBUG);
+
+  sfd = do_connect();
+  if (sfd < 0) {
+    return EXIT_FAILURE;
+  }
+
+  aes_key = get_aes_key(sfd, &key_len);
+  if (aes_key == NULL) {
+    close(sfd);
+    return EXIT_FAILURE;
+  }
+
+  fractions = fetch_fractions(sfd, &fraction_count);
+  if (fractions == NULL) {
+    free(aes_key);
+    close(sfd);
+    return EXIT_FAILURE;
+  }
+
+  log_info("Downloaded fractions");
+
+  qsort(fractions, fraction_count, sizeof(fraction_t), compare_fractions);
+
+  module = decrypt_lkm(fractions, fraction_count, &module_size, aes_key);
   if (module == NULL) {
     log_error("There was an error creating the module");
-    goto cleanup;
+    cleanup_fraction_array(fractions, fraction_count);
+    free(aes_key);
+    close(sfd);
+    return EXIT_FAILURE;
   }
 
   if (load_lkm(module, module_size) < 0) {
     log_error("Error loading LKM");
-    goto cleanup;
+    free(module);
+    cleanup_fraction_array(fractions, fraction_count);
+    free(aes_key);
+    close(sfd);
+    return EXIT_FAILURE;
   }
 
-  http_free(&http_post_res);
-  http_free(&http_fraction_res);
-  cleanup_string_array(fraction_links, num_links);
-  cleanup_fraction_array(fractions, lines_read);
   free(module);
-  free(private_key);
-  free(public_key);
-  EVP_PKEY_free(pkey);
-
+  cleanup_fraction_array(fractions, fraction_count);
+  free(aes_key);
   close(sfd);
 
   return EXIT_SUCCESS;
-
-  /* There's nothing to see here, move on*/
-cleanup: // we accept NO comments on this. have a !nice day
-  if (sfd != -1) {
-    close(sfd);
-  }
-  if (fraction_links) {
-    cleanup_string_array(fraction_links, num_links);
-  }
-  if (fractions) {
-    cleanup_fraction_array(fractions, num_links);
-  }
-  if (http_fraction_res.data) {
-    http_free(&http_fraction_res);
-  }
-  if (http_post_res.data) {
-    http_free(&http_post_res);
-  }
-  free(module); // no need to check module != NULL as free(NULL) is defined by
-  // the C standard to do nothing
-  if (private_key) {
-    free(private_key);
-  }
-  if (public_key) {
-    free(public_key);
-  }
-  if (pkey) {
-    EVP_PKEY_free(pkey);
-  }
-
-  return EXIT_FAILURE;
 }

client/src/utils.c

@@ -4,29 +4,6 @@
 #include <stdlib.h>
 #include <time.h>
 
-// if this dont work Skelly is to blame ;)
-int split_fraction_links(char *data, char *data_arr[], int maxlines) {
-  int lines_read = 0;
-  char *line;
-
-  // Use the input `data` directly, no need for strdup
-  line = strtok(data, "\n");
-  while (line != NULL && lines_read < maxlines) {
-    data_arr[lines_read] = strdup(line);
-    if (data_arr[lines_read] == NULL) {
-      log_error("strdup failed to allocate memory");
-      // Free previously allocated lines in case of failure
-      for (int i = 0; i < lines_read; i++) {
-        free(data_arr[i]);
-      }
-      return -1;
-    }
-    lines_read++;
-    line = strtok(NULL, "\n");
-  }
-  return lines_read;
-}
-
 void print_hex(const unsigned char *data, size_t size) {
   if (data == NULL) {
     log_error("Null data pointer\n"); // had to learn the hard f way