8 files changed, 197 insertions, 21 deletions

diff --git a/src/wav.c b/src/wav.c
index 61765d6..665998d 100644
--- a/src/wav.c
+++ b/src/wav.c
@@ -5,48 +5,90 @@
 #include <stdio.h>
 #include <string.h>
 
-void seek_to_chunk(FILE *file, char *chunk_name) {
+b32 seek_to_chunk(FILE *file, char const *chunk_name) {
     char cur_name[4];
      u32  chunk_size;
 
+    fseek(file, 0, SEEK_END);
+    const i64 max_pos = ftell(file);
     fseek(file, sizeof(wav_master_chunk_t), SEEK_SET);
     while (true) {
         fread(cur_name, 4, 1, file);
         fread(&chunk_size, 4, 1, file);
 
-        if (strcmp(cur_name, chunk_name) == 0) {
-            break;
-        }
+        if (memcmp(cur_name, chunk_name, 4) == 0) break;
+
+        if (chunk_size % 2 == 1) chunk_size++;
+
+        if (ftell(file) + (i64)chunk_size >= max_pos) return false;
 
         fseek(file, chunk_size, SEEK_CUR);
     }
 
     fseek(file, -8, SEEK_CUR);
+    return true;
 }
 
-wav_data_t *load_wav_file(smrt_arena_t *arena, char *filename, wav_master_chunk_t *master_o, wav_fmt_chunk_t *format_o) {
+wav_data_t load_wav_file(smrt_arena_t *arena, char *filename, wav_master_chunk_t *master_o, wav_fmt_chunk_t *format_o) {
+    wav_data_t data = { 0 };
+
     FILE *wav_file = fopen(filename, "rb");
 
-    if (!wav_file) {
-        return NULL;
-    }
+    if (!wav_file) return data;
 
     fread(master_o, sizeof(wav_master_chunk_t), 1, wav_file);
 
-    seek_to_chunk(wav_file, "fmt ");
+    if (seek_to_chunk(wav_file, "fmt ") == false) return data;
     fread(format_o, sizeof(wav_fmt_chunk_t), 1, wav_file);
 
-    seek_to_chunk(wav_file, "data");
+    if (seek_to_chunk(wav_file, "data") == false) return data;
     fseek(wav_file, 4, SEEK_CUR);
+
     u32 sampled_data_size;
     fread(&sampled_data_size, 4, 1, wav_file);
 
-    wav_data_t *data = smrt_arena_push(arena, sizeof(wav_data_t) + sampled_data_size, true);
-    fread((u8*)data+WAV_DATA_BASE_POS, sampled_data_size, 1, wav_file);
-    data->sample_count = sampled_data_size / ((format_o->bits_per_sample / 8) * format_o->num_channels);
+    data.samples = smrt_arena_push(arena, sizeof(wav_data_t) + sampled_data_size, true);
+    if (!data.samples) return data;
+
+    fread((u8*)data.samples+WAV_DATA_BASE_POS, sampled_data_size, 1, wav_file);
+    data.sample_count = sampled_data_size / ((format_o->bits_per_sample / 8) * format_o->num_channels);
 
     return data;
 }
 
-void write_wav_file(char *filename, wav_fmt_chunk_t *header_i, wav_data_t *data_i);
-wav_fmt_chunk_t make_wav_fmt_chunk(u32 num_channels, u32 sample_rate, u16 bits_per_sample, wav_data_t *data_i);
+b32 write_wav_file(char *filename, wav_fmt_chunk_t *fmt_chunk_i, wav_data_t data_i) {
+    FILE *wav_file = fopen(filename, "wb");
+
+    if (!wav_file) return false;
+
+    u64 data_size_bytes = data_i.sample_count * fmt_chunk_i->bytes_per_block;
+
+    wav_master_chunk_t mchunk = {
+        .file_size = sizeof(wav_master_chunk_t) + sizeof(wav_fmt_chunk_t) + sizeof(wav_data_t) + data_size_bytes
+    };
+
+    memcpy(&mchunk.chunk_id, "RIFF", 4);
+    memcpy(&mchunk.file_format_id, "WAVE", 4);
+
+    fwrite(&mchunk, sizeof(wav_master_chunk_t), 1, wav_file);
+    fwrite(fmt_chunk_i, sizeof(wav_data_t), 1, wav_file);
+    fwrite(data_i.samples, sizeof(u8), data_size_bytes, wav_file);
+
+    return true;
+}
+
+wav_fmt_chunk_t make_wav_fmt_chunk(u32 num_channels, u32 sample_rate, u16 bits_per_sample, wav_data_t data_i) {
+    u64 data_size_bytes = data_i.sample_count * (num_channels * (bits_per_sample / 8));
+    wav_fmt_chunk_t o = {
+        .chunk_size = sizeof(wav_fmt_chunk_t) + data_size_bytes - 8,
+        .bits_per_sample = bits_per_sample,
+        .num_channels = num_channels,
+        .sample_rate = sample_rate,
+        .bytes_per_block = num_channels * (bits_per_sample / 8),
+        .audio_format = 1,
+    };
+    o.bytes_per_sec = o.bytes_per_block * sample_rate;
+    memcpy(&o.chunk_id, "data", 4);
+
+    return o;
+}
diff --git a/src/wav.h b/src/wav.h
index dd85ccd..773cc2a 100644
--- a/src/wav.h
+++ b/src/wav.h
@@ -7,6 +7,7 @@
 
 typedef struct {
     u64    sample_count;
+    u8    *samples;
 } wav_data_t;
 
 // [Master RIFF chunk]
@@ -35,12 +36,12 @@ typedef struct {
     u16    audio_format;
     u16    num_channels;
     u32     sample_rate;
-    u32   byter_per_sec;
+    u32   bytes_per_sec;
     u16 bytes_per_block;
     u16 bits_per_sample;
 } wav_fmt_chunk_t;
 #pragma pack(pop)
 
-     wav_data_t     *load_wav_file(smrt_arena_t *arena, char *filename, wav_master_chunk_t *master_o, wav_fmt_chunk_t *format_o);
-           void     write_wav_file(char *filename, wav_fmt_chunk_t *header_i, wav_data_t *data_i);
-wav_fmt_chunk_t make_wav_fmt_chunk(u32 num_channels, u32 sample_rate, u16 bits_per_sample, wav_data_t *data_i);
+     wav_data_t      load_wav_file(smrt_arena_t *arena, char *filename, wav_master_chunk_t *master_o, wav_fmt_chunk_t *format_o);
+            b32     write_wav_file(char *filename, wav_fmt_chunk_t *fmt_chunk_i, wav_data_t data_i);
+wav_fmt_chunk_t make_wav_fmt_chunk(u32 num_channels, u32 sample_rate, u16 bits_per_sample, wav_data_t data_i);
diff --git a/tests/fixtures/extra_chunk.wav b/tests/fixtures/extra_chunk.wav
new file mode 100644
index 0000000..21e346c
--- /dev/null
+++ b/tests/fixtures/extra_chunk.wav
diff --git a/tests/fixtures/missing_data_chunk.wav b/tests/fixtures/missing_data_chunk.wav
new file mode 100644
index 0000000..89e1533
--- /dev/null
+++ b/tests/fixtures/missing_data_chunk.wav
diff --git a/tests/fixtures/mono8.wav b/tests/fixtures/mono8.wav
new file mode 100644
index 0000000..6c8fe8a
--- /dev/null
+++ b/tests/fixtures/mono8.wav
diff --git a/tests/fixtures/odd_chunk.wav b/tests/fixtures/odd_chunk.wav
new file mode 100644
index 0000000..a1f1274
--- /dev/null
+++ b/tests/fixtures/odd_chunk.wav
diff --git a/tests/fixtures/stereo16.wav b/tests/fixtures/stereo16.wav
new file mode 100644
index 0000000..53795c8
--- /dev/null
+++ b/tests/fixtures/stereo16.wav
diff --git a/tests/wav.c b/tests/wav.c
index aac7f9d..81b860e 100644
--- a/tests/wav.c
+++ b/tests/wav.c
@@ -14,14 +14,147 @@ Test(wav, read_file) {
     wav_master_chunk_t mchunk;
     wav_fmt_chunk_t  fmtchunk;
 
-    wav_data_t *data = load_wav_file(arena, "./tests/fixtures/sine.wav", &mchunk, &fmtchunk);
+    wav_data_t data = load_wav_file(arena, "./tests/fixtures/sine.wav", &mchunk, &fmtchunk);
 
     u8 idcheck[4] = { 'f', 'm', 't', ' ' };
     cr_expect(memcmp(fmtchunk.chunk_id, idcheck, 4) == 0);
     cr_expect(fmtchunk.audio_format == 1);
     cr_expect(fmtchunk.num_channels == 1);
     cr_expect(fmtchunk.bits_per_sample == 16);
-    cr_expect(data->sample_count == 48000);
+    cr_expect(data.sample_count == 48000);
+
+    smrt_arena_destroy(arena);
+}
+
+Test(wav, sample_data_matches_source) {
+    smrt_arena_t *arena = smrt_arena_create(MiB(4), KiB(4), false);
+
+    wav_master_chunk_t mchunk;
+    wav_fmt_chunk_t  fmtchunk;
+
+    wav_data_t data = load_wav_file(arena, "./tests/fixtures/sine.wav", &mchunk, &fmtchunk);
+
+    i16 expected[8] = { 0, 2147, 4242, 6291, 8185, 9935, 11580, 12979 };
+    i16 *samples = (i16*)((u8*)data.samples + WAV_DATA_BASE_POS);
+
+    for (int i = 0; i < 8; i++) {
+        cr_expect_eq(samples[i], expected[i]);
+    }
+
+    smrt_arena_destroy(arena);
+}
+
+Test(wav, missing_file_returns_null) {
+    smrt_arena_t *arena = smrt_arena_create(MiB(4), KiB(4), false);
+
+    wav_master_chunk_t mchunk;
+    wav_fmt_chunk_t  fmtchunk;
+
+    wav_data_t data = load_wav_file(arena, "./tests/fixtures/does_not_exist.wav", &mchunk, &fmtchunk);
+
+    cr_expect_eq(data.samples, NULL);
+
+    smrt_arena_destroy(arena);
+}
+
+Test(wav, stereo_16bit) {
+    smrt_arena_t *arena = smrt_arena_create(MiB(4), KiB(4), false);
+
+    wav_master_chunk_t mchunk;
+    wav_fmt_chunk_t  fmtchunk;
+
+    wav_data_t data = load_wav_file(arena, "./tests/fixtures/stereo16.wav", &mchunk, &fmtchunk);
+
+    cr_expect(fmtchunk.num_channels == 2);
+    cr_expect(fmtchunk.bits_per_sample == 16);
+    cr_expect(data.sample_count == 5);
+
+    i16 *samples = (i16*)((u8*)data.samples + WAV_DATA_BASE_POS);
+    cr_expect_eq(samples[0], 0);
+    cr_expect_eq(samples[1], 0);
+    cr_expect_eq(samples[2], 100);
+    cr_expect_eq(samples[3], -100);
+
+    smrt_arena_destroy(arena);
+}
+
+Test(wav, mono_8bit) {
+    smrt_arena_t *arena = smrt_arena_create(MiB(4), KiB(4), false);
+
+    wav_master_chunk_t mchunk;
+    wav_fmt_chunk_t  fmtchunk;
+
+    wav_data_t data = load_wav_file(arena, "./tests/fixtures/mono8.wav", &mchunk, &fmtchunk);
+
+    cr_expect(fmtchunk.num_channels == 1);
+    cr_expect(fmtchunk.bits_per_sample == 8);
+    cr_expect(data.sample_count == 6);
+
+    u8 *samples = (u8*)data.samples + WAV_DATA_BASE_POS;
+    u8 expected[6] = { 0, 64, 128, 192, 255, 32 };
+    for (int i = 0; i < 6; i++) {
+        cr_expect_eq(samples[i], expected[i]);
+    }
+
+    smrt_arena_destroy(arena);
+}
+
+Test(wav, skips_unknown_chunk_before_fmt) {
+    smrt_arena_t *arena = smrt_arena_create(MiB(4), KiB(4), false);
+
+    wav_master_chunk_t mchunk;
+    wav_fmt_chunk_t  fmtchunk;
+
+    wav_data_t data = load_wav_file(arena, "./tests/fixtures/extra_chunk.wav", &mchunk, &fmtchunk);
+
+    u8 idcheck[4] = { 'f', 'm', 't', ' ' };
+    cr_expect(memcmp(fmtchunk.chunk_id, idcheck, 4) == 0);
+    cr_expect(fmtchunk.num_channels == 1);
+    cr_expect(fmtchunk.sample_rate == 22050);
+    cr_expect(data.sample_count == 4);
+
+    smrt_arena_destroy(arena);
+}
+
+Test(wav, odd_sized_chunk_before_fmt, .timeout = 2) {
+    smrt_arena_t *arena = smrt_arena_create(MiB(4), KiB(4), false);
+
+    wav_master_chunk_t mchunk;
+    wav_fmt_chunk_t  fmtchunk;
+
+    wav_data_t data = load_wav_file(arena, "./tests/fixtures/odd_chunk.wav", &mchunk, &fmtchunk);
+
+    u8 idcheck[4] = { 'f', 'm', 't', ' ' };
+    cr_expect(memcmp(fmtchunk.chunk_id, idcheck, 4) == 0);
+    cr_expect(fmtchunk.sample_rate == 16000);
+    cr_expect(data.sample_count == 4);
+
+    smrt_arena_destroy(arena);
+}
+
+Test(wav, missing_data_chunk, .timeout = 2) {
+    smrt_arena_t *arena = smrt_arena_create(MiB(4), KiB(4), false);
+
+    wav_master_chunk_t mchunk;
+    wav_fmt_chunk_t  fmtchunk;
+
+    wav_data_t data = load_wav_file(arena, "./tests/fixtures/missing_data_chunk.wav", &mchunk, &fmtchunk);
+
+    cr_expect_eq(data.samples, NULL);
+
+    smrt_arena_destroy(arena);
+}
+
+Test(wav, arena_too_small_for_data) {
+    smrt_arena_t *arena = smrt_arena_create(KiB(1), KiB(4), false);
+
+    wav_master_chunk_t mchunk;
+    wav_fmt_chunk_t  fmtchunk;
+
+    wav_data_t data = load_wav_file(arena, "./tests/fixtures/sine.wav", &mchunk, &fmtchunk);
+
+    cr_expect_eq(data.samples, NULL);
+    cr_expect_eq(data.sample_count, 0);
 
     smrt_arena_destroy(arena);
 }