Fixed memory leaking issues and introduced a feature to list the contents of an archive.

9 files changed, 135 insertions, 45 deletions

diff --git a/src/cli.c b/src/cli.c
index 6144c4e..11638a2 100644
--- a/src/cli.c
+++ b/src/cli.c
@@ -31,7 +31,9 @@
                   "exit.\n" \
                   "   -v, --version\tDisplay the currently installed " \
                   "version and exit.\n\n" \
-                  "   -l, --list\t\tList games supported by Nekopack.\n" \
+                  "   -e, --extract\tExtract the contents of the archive. " \
+                  "This is the default.\n" \
+                  "   -l, --list\t\tList the contents of the archive.\n\n" \
                   "   -a, --archive\tAlternate way of specifying archive" \
                   "to extract.\n" \
                   "   -g, --game\t\tGame the archive is from. Required for " \
@@ -64,36 +66,43 @@ struct configuration parse_args(int argc, char *argv[]) {
     };
 
     do {
-        current = getopt_long(argc, argv, "hvla:g:", long_options,
+        current = getopt_long(argc, argv, "hvela:g:", long_options,
                               &option_index);
         switch (current) {
             case 'h':
                 printf("Usage: %s [OPTIONS] (ARCHIVE PATH)\n\n", argv[0]);
-                printf("%s\n", HELP_TEXT);
+                printf("%s\n\n", HELP_TEXT);
+                printf("Supported Game Crypto:\n%s\n", GAME_CONSTANTS);
                 exit(EXIT_SUCCESS);
             case 'v':
                 printf("Nekopack, version %s\nProgrammed by "
                        "Jakob. <http://tsar-fox.com/>\n", VERSION);
                 exit(EXIT_SUCCESS);
-            case 'l':
-                printf("%s\n", GAME_CONSTANTS);
-                exit(EXIT_SUCCESS);
             case 'a':
                 count++;
                 parsed.archive_path = optarg;
                 break;
             case 'g':
                 count++;
+                /* Out of concern for efficiency, string comparisons are
+                   done here and the source game is stored in an enum. */
                 if (!strcmp(optarg, "nekopara_volume_0"))
-                    parsed.source = NEKOPARA_VOLUME_0;
+                    parsed.game = NEKOPARA_VOLUME_0;
                 else if (!strcmp(optarg, "nekopara_volume_0_steam"))
-                    parsed.source = NEKOPARA_VOLUME_0_STEAM;
+                    parsed.game = NEKOPARA_VOLUME_0_STEAM;
                 else if (!strcmp(optarg, "nekopara_volume_1"))
-                    parsed.source = NEKOPARA_VOLUME_1;
+                    parsed.game = NEKOPARA_VOLUME_1;
                 else if (!strcmp(optarg, "nekopara_volume_1_steam"))
-                    parsed.source = NEKOPARA_VOLUME_1_STEAM;
+                    parsed.game = NEKOPARA_VOLUME_1_STEAM;
                 else
-                    parsed.source = NO_CRYPTO;
+                    parsed.game = NO_CRYPTO;
+                break;
+            case 'e':
+                /* '-e' is still parsed to provide consisency. */
+                parsed.mode = EXTRACT;
+                break;
+            case 'l':
+                parsed.mode = LIST;
         }
         count++;
     } while (current >= 0);
diff --git a/src/cli.h b/src/cli.h
index 43d9465..2b2d988 100644
--- a/src/cli.h
+++ b/src/cli.h
@@ -20,17 +20,25 @@
 /* Enumerable type representing Nekopara games, as they have different
    encryption keys. It's stored as an enum because multiple strcmp calls
    to figure out how to decrypt is a waste of processor cycles. */
-typedef enum game {
+typedef enum game_type {
     NO_CRYPTO,
     NEKOPARA_VOLUME_0,
     NEKOPARA_VOLUME_0_STEAM,
     NEKOPARA_VOLUME_1,
     NEKOPARA_VOLUME_1_STEAM,
-} game;
+} game_type;
+
+/* Enumerable type for mode of operation. Used in the main function
+   to decide what to do after the initial XP3 sanity checks. */
+typedef enum mode_type {
+    EXTRACT,
+    LIST,
+} mode_type;
 
 /* Binary structure for storing command-line options. */
 struct configuration {
-    game source; /* Which decryption key to use. */
+    game_type game; /* Which decryption key to use. */
+    mode_type mode; /* What to do after initial sanity checks. */
     const char *archive_path; /* Path to archive to extract. */
 };
 
diff --git a/src/crypto.c b/src/crypto.c
index 2d527c7..f6e2cd8 100644
--- a/src/crypto.c
+++ b/src/crypto.c
@@ -20,9 +20,9 @@
 
 
 /* Returns the encryption keys for a given game value. */
-key get_encryption_key(game current_game) {
+key get_encryption_key(game_type game) {
     key encryption_key;
-    switch (current_game) {
+    switch (game) {
         case NEKOPARA_VOLUME_0:
             encryption_key.master_key = 0x1548e29c;
             encryption_key.initial_fallback_key = 0x9c;
diff --git a/src/crypto.h b/src/crypto.h
index e82b934..f8d5d86 100644
--- a/src/crypto.h
+++ b/src/crypto.h
@@ -32,7 +32,7 @@ typedef struct {
 } key;
 
 /* Returns the encryption keys for a given game value. */
-key get_encryption_key(game current_game);
+key get_encryption_key(game_type game);
 
 /* Decrypts the contents of a buffer according to a file key. */
 void decrypt_buffer(Bytef *encrypted_buffer, uint64_t buffer_length,
diff --git a/src/extract.c b/src/extract.c
index 3b6828b..261d84a 100644
--- a/src/extract.c
+++ b/src/extract.c
@@ -15,7 +15,6 @@
    You should have received a copy of the GNU General Public License
    along with Nekopack. If not, see <http://www.gnu.org/licenses/>. */
 
-#include <inttypes.h> // Needed for debugging at this point.
 #include <stdio.h>
 #include <stdint.h>
 #include <stdlib.h>
@@ -26,7 +25,6 @@
 
 #include "cli.h"
 #include "defs.h"
-#include "decompress.h"
 #include "extract.h"
 #include "file.h"
 #include "write.h"
@@ -40,9 +38,9 @@ elif_node *read_elif_entry(memory_stream *data_stream);
 void read_stream(void *destination, Bytef **source, size_t size);
 
 
-/* Handles decompression of the archive, as well as
-   parsing, decrypting and writing the table entries. */
-void extract(FILE *archive, uint64_t table_offset) {
+/* Decrypts and writes files in the XP3 archive
+   to disk according to table entries. */
+void extract(memory_stream data_stream, FILE *archive) {
     /* eliF and File entries are stored in a linked list as they're
        seen because the order of entries in XP3 archives is not
        guaranteed to be chronological. calloc is used to prevent the
@@ -50,23 +48,7 @@ void extract(FILE *archive, uint64_t table_offset) {
     elif_node *elif_new, *elif_root = calloc(sizeof(elif_node), 1);
     file_node *file_new, *file_root = calloc(sizeof(file_node), 1);
 
-    uint8_t compressed;
-    uint64_t compressed_size, decompressed_size;
-    fseek(archive, table_offset, SEEK_SET);
-    fread(&compressed, sizeof(uint8_t), 1, archive);
-    fread(&compressed_size, sizeof(uint64_t), 1, archive);
-    fread(&decompressed_size, sizeof(uint64_t), 1, archive);
-
     int stream_ended = 0;
-    memory_stream data_stream;
-    memory_stream compressed_data = read_to_stream(archive, compressed_size);
-    if (compressed) {
-        data_stream = decompress_stream(compressed_data, decompressed_size);
-        free(compressed_data.start);
-    } else {
-        data_stream = compressed_data;
-    }
-
     uint32_t entry_magic;
     uint64_t entry_size;
     do {
@@ -82,11 +64,23 @@ void extract(FILE *archive, uint64_t table_offset) {
             case HNFN_MAGIC:
             case NEKO_MAGIC:
                 elif_new = read_elif_entry(&data_stream);
+                if (elif_new == NULL) {
+                    fprintf(stderr, "Insufficient memory.\n");
+                    free_elif_nodes(elif_root);
+                    free_file_nodes(file_root);
+                    return;
+                }
                 defer_elif_node(elif_new, elif_root);
                 break;
             case FILE_MAGIC:
                 file_new = read_file_entry(&data_stream,
                                            data_stream.data + entry_size);
+                if (file_new == NULL) {
+                    fprintf(stderr, "Insufficient memory.\n");
+                    free_elif_nodes(elif_root);
+                    free_file_nodes(file_root);
+                    return;
+                }
                 defer_file_node(file_new, file_root);
                 break;
             default:
@@ -97,7 +91,35 @@ void extract(FILE *archive, uint64_t table_offset) {
     write_files(file_root, elif_root, archive);
     free_elif_nodes(elif_root);
     free_file_nodes(file_root);
-    free(data_stream.start);
+}
+
+
+/* Simply lists the contents of an archive, ignoring File entries. */
+void list(memory_stream data_stream) {
+    int stream_ended = 0;
+    uint32_t entry_magic;
+    uint64_t entry_size;
+    elif_node *temporary;
+    do {
+        read_stream(&entry_magic, &data_stream.data, sizeof(uint32_t));
+        read_stream(&entry_size, &data_stream.data, sizeof(uint64_t));
+
+        switch (entry_magic) {
+            case ELIF_MAGIC:
+            case HNFN_MAGIC:
+            case NEKO_MAGIC:
+                temporary = read_elif_entry(&data_stream);
+                printf("%s\n", temporary->file_name);
+                free(temporary->file_name);
+                free(temporary);
+                break;
+            case FILE_MAGIC:
+                data_stream.data += entry_size;
+                break;
+            default:
+                stream_ended = 1;
+        }
+    } while (!stream_ended);
 }
 
 
@@ -123,12 +145,18 @@ elif_node *read_elif_entry(memory_stream *data_stream) {
            which aren't counted in the name size. */
         char *input_buffer = malloc(name_size * 2 + 2);
         read_stream(input_buffer, &data_stream->data, name_size * 2 + 2);
+        if (input_buffer == NULL)
+            return NULL;
 
         /* name_size + 1 is unreliable since we're going back to UTF-8
            and some characters (especially Japanese ones) can be wide.
            To compensate we just allocate a buffer that could fit the
            UTF-16LE data and utilize null-bytes to terminate names. */
         file_name = malloc(name_size * 2 + 2);
+        if (file_name == NULL) {
+            free(input_buffer);
+            return NULL;
+        }
 
         /* iconv is the less-portable glibc way of doing it. It seems to
            be in the OpenBSD manpages, though, so I'm not worried. */
@@ -146,6 +174,10 @@ elif_node *read_elif_entry(memory_stream *data_stream) {
     }
 
     elif_node *current = malloc(sizeof(elif_node));
+    if (current == NULL) {
+        free(file_name);
+        return NULL;
+    }
     current->key = file_key;
     current->file_name = file_name;
     current->next = NULL;
diff --git a/src/extract.h b/src/extract.h
index 9aea965..10965a5 100644
--- a/src/extract.h
+++ b/src/extract.h
@@ -30,9 +30,12 @@ typedef struct {
     Bytef *data;
 } memory_stream;
 
-/* Handles decompression of the archive, as well as
-   parsing, decrypting and writing the table entries. */
-void extract(FILE *archive, uint64_t table_offset);
+/* Decrypts and writes files in the XP3 archive
+   to disk according to table entries. */
+void extract(memory_stream data_stream, FILE *archive);
+
+/* Simply lists the contents of an archive, ignoring File entries. */
+void list(memory_stream data_stream);
 
 /* Wrapper for memcpy which increments the source operand by
    the amount of bytes read to simulate a file stream. */
diff --git a/src/file.c b/src/file.c
index 12e5432..3274859 100644
--- a/src/file.c
+++ b/src/file.c
@@ -40,6 +40,8 @@ file_node *read_file_entry(memory_stream *data_stream, Bytef *section_end) {
     uint32_t entry_magic;
     uint64_t entry_size;
     file_node *parsed = calloc(sizeof(file_node), 1);
+    if (parsed == NULL)
+        return NULL;
     while (data_stream->data < section_end) {
         read_stream(&entry_magic, &data_stream->data, sizeof(uint32_t));
         read_stream(&entry_size, &data_stream->data, sizeof(uint64_t));
@@ -50,6 +52,11 @@ file_node *read_file_entry(memory_stream *data_stream, Bytef *section_end) {
             case SEGM_MAGIC:
                 /* Segments are 28 bytes each. */
                 read_segm_chunk(data_stream, parsed, entry_size / 28);
+                /* Check if segments was successfully allocated. */
+                if (parsed->segments == NULL) {
+                    free(parsed);
+                    return NULL;
+                }
                 break;
             case INFO_MAGIC:
                 read_info_chunk(data_stream, parsed);
@@ -78,6 +85,7 @@ void read_info_chunk(memory_stream *data_stream, file_node *parsed) {
 
     char *file_name = malloc(file_name_size * 2 + 2);
     read_stream(file_name, &data_stream->data, file_name_size * 2 + 2);
+    free(file_name);
 }
 
 
@@ -86,6 +94,8 @@ void read_segm_chunk(memory_stream *data_stream, file_node *parsed,
                      uint64_t segment_count) {
     /* Segments are stored in an array of segment pointers. */
     segment **segments = malloc(sizeof(segment *) * segment_count);
+    if (segments == NULL)
+        return;
     for (uint64_t i = 0; i < segment_count; i++) {
         segments[i] = malloc(sizeof(segment));
         read_stream(&segments[i]->compressed,
diff --git a/src/main.c b/src/main.c
index 662acee..36ddc01 100644
--- a/src/main.c
+++ b/src/main.c
@@ -22,6 +22,7 @@
 #include <string.h>
 
 #include "cli.h"
+#include "decompress.h"
 #include "extract.h"
 
 #define XP3_MAGIC "XP3\x0d\x0a\x20\x0a\x1a\x8b\x67\x01"
@@ -51,7 +52,34 @@ int main(int argc, char *argv[]) {
 
     int archive_version = get_archive_version(archive);
     uint64_t table_offset = get_table_offset(archive, archive_version);
-    extract(archive, table_offset);
+
+    uint8_t compressed;
+    uint64_t compressed_size, decompressed_size;
+    fseek(archive, table_offset, SEEK_SET);
+    fread(&compressed, sizeof(uint8_t), 1, archive);
+    fread(&compressed_size, sizeof(uint64_t), 1, archive);
+    fread(&decompressed_size, sizeof(uint64_t), 1, archive);
+
+    /* Every task that hasn't already been handled needs a decompressed
+       instance of the archive table. Decompression done in memory
+       because it's $CURRENT_YEAR. */
+    memory_stream data_stream;
+    memory_stream compressed_data = read_to_stream(archive, compressed_size);
+    if (compressed) {
+        data_stream = decompress_stream(compressed_data, decompressed_size);
+        free(compressed_data.start);
+    } else {
+        data_stream = compressed_data;
+    }
+
+    switch (arguments.mode) {
+        case EXTRACT:
+            extract(data_stream, archive);
+            break;
+        case LIST:
+            list(data_stream);
+    }
+    free(data_stream.start);
     fclose(archive);
     return 0;
 }
diff --git a/src/write.c b/src/write.c
index e011504..ebf640c 100644
--- a/src/write.c
+++ b/src/write.c
@@ -55,7 +55,7 @@ char *pop_file_name(uint32_t key, elif_node *root) {
    entry to disk, according to information specified by the node. */
 void write_files(file_node *file_root, elif_node *elif_root, FILE *archive) {
     file_node *current;
-    key encryption_key = get_encryption_key(arguments.source);
+    key encryption_key = get_encryption_key(arguments.game);
     Bytef *compressed_buffer, *decompressed_buffer, *out_buffer, *out_start;
     for (current = file_root->next; current != NULL; current = current->next) {
         char *file_name = pop_file_name(current->key, elif_root);
@@ -84,7 +84,7 @@ void write_files(file_node *file_root, elif_node *elif_root, FILE *archive) {
             free(decompressed_buffer);
         }
 
-        if (arguments.source != NO_CRYPTO) {
+        if (arguments.game != NO_CRYPTO) {
             decrypt_buffer(out_start, current->file_size,
                            encryption_key, current->key);
         }