// Test suite for Secure_Memory.
//
// module.jai #loads this file when the module is imported with LOAD_TESTS = true, so the
// suite is built against the same USE_SODIUM setting the module was built with -- running it
// tells you the wipe path you actually shipped works, not the default one.
//
// Every check writes a recognizable pattern, wipes part of it, and asserts both that the
// wiped span changed and that the bytes on either side did not.

#import "Basic";   // String_Builder, print, print_to_builder, tprint

#scope_file

MAX_GROUPS :: 8;

group_names  : [MAX_GROUPS] string;
group_checks : [MAX_GROUPS] int;
group_fails  : [MAX_GROUPS] int;
group_count  := 0;

checks_run    := 0;
checks_failed := 0;

// Set by run_secure_memory_tests. One suite is one work item, so a file-scope pointer is
// enough; null means write straight to stdout, as a synchronous or standalone run does.
output_builder: *String_Builder;

emit :: (format_string: string, args: ..Any) {
  if output_builder  print_to_builder(output_builder, format_string, ..args);
  else               print(format_string, ..args);
} @PrintLike

group :: (name: string) {
  assert(group_count < MAX_GROUPS, "raise MAX_GROUPS");
  group_names[group_count] = name;
  group_count += 1;
  emit("\n-- % --\n", name);
}

expect :: (condition: bool, format: string, args: ..Any, loc := #caller_location) {
  checks_run += 1;
  if group_count  group_checks[group_count-1] += 1;
  if condition return;

  checks_failed += 1;
  if group_count  group_fails[group_count-1] += 1;

  name := ifx group_count then group_names[group_count-1] else "?";
  emit("  FAIL  [%:%]  %\n", name, loc.line_number, tprint(format, ..args));
}

fill :: (memory: *void, count: s64, value: u8) {
  bytes := cast(*u8) memory;
  for 0..count-1  bytes[it] = value;
}

all_bytes_are :: (memory: *void, count: s64, value: u8) -> bool {
  bytes := cast(*u8) memory;
  for 0..count-1  if bytes[it] != value  return false;
  return true;
}

// A buffer with a canary run on each side, so a wipe that walks past its span is visible.
GUARD  :: 16;
BODY   :: 64;
CANARY :: 0xa5;

Guarded :: struct {
  bytes: [GUARD + BODY + GUARD] u8;
}

body_of :: (using g: *Guarded) -> *u8 {
  return bytes.data + GUARD;
}

prepare :: (g: *Guarded, pattern: u8) {
  fill(g.bytes.data, g.bytes.count, CANARY);
  fill(body_of(g), BODY, pattern);
}

guards_intact :: (g: *Guarded) -> bool {
  if !all_bytes_are(g.bytes.data, GUARD, CANARY)                  return false;
  if !all_bytes_are(g.bytes.data + GUARD + BODY, GUARD, CANARY)   return false;
  return true;
}

// ---------------------------------------------------------------------------

memory_secure_overwrite_tests :: () {
  group("memory_secure_overwrite");

  g: Guarded;

  prepare(*g, 0xcc);
  memory_secure_overwrite(body_of(*g), BODY);
  expect(all_bytes_are(body_of(*g), BODY, 0), "default value did not zero the body");
  expect(guards_intact(*g), "default value ran past the body");

  prepare(*g, 0xcc);
  memory_secure_overwrite(body_of(*g), BODY, 0xff);
  expect(all_bytes_are(body_of(*g), BODY, 0xff), "explicit fill value was not written");
  expect(guards_intact(*g), "explicit fill value ran past the body");

  // A partial wipe must leave the tail of the body alone, not just the guards.
  HALF :: BODY / 2;
  prepare(*g, 0xcc);
  memory_secure_overwrite(body_of(*g), HALF);
  expect(all_bytes_are(body_of(*g), HALF, 0), "partial wipe did not zero the first half");
  expect(all_bytes_are(body_of(*g) + HALF, BODY - HALF, 0xcc), "partial wipe clobbered the second half");
  expect(guards_intact(*g), "partial wipe ran past the body");

  prepare(*g, 0xcc);
  memory_secure_overwrite(body_of(*g), 0);
  expect(all_bytes_are(body_of(*g), BODY, 0xcc), "a zero count wrote something");
  expect(guards_intact(*g), "a zero count ran past the body");
}

secure_overwrite_tests :: () {
  group("secure_overwrite overloads");

  // string: count is already a byte count.
  {
    storage: [BODY] u8;
    fill(storage.data, storage.count, 0xcc);
    s := string.{storage.count, storage.data};
    secure_overwrite(s);
    expect(all_bytes_are(storage.data, storage.count, 0), "string overload did not zero the backing bytes");
  }

  // []$T: the overload has to scale count by size_of(T), or it wipes a quarter of the array.
  {
    values: [16] u32;
    for * values  it.* = 0xdead_beef;
    view: []u32 = values;
    secure_overwrite(view);
    expect(all_bytes_are(values.data, values.count * size_of(u32), 0), "slice overload did not zero every element");
  }

  // *$T: wipes exactly size_of(T).
  {
    Blob :: struct {
      a: u64;
      b: u64;
      c: u64;
    }
    pair: [2] Blob;
    fill(pair.data, size_of(Blob) * 2, 0xcc);
    secure_overwrite(*pair[0]);
    expect(all_bytes_are(*pair[0], size_of(Blob), 0), "pointer overload did not zero the struct");
    expect(all_bytes_are(*pair[1], size_of(Blob), 0xcc), "pointer overload ran into the next struct");
  }

  // *$T + explicit byte count.
  {
    g: Guarded;
    prepare(*g, 0xcc);
    secure_overwrite(body_of(*g), BODY);
    expect(all_bytes_are(body_of(*g), BODY, 0), "counted overload did not zero the body");
    expect(guards_intact(*g), "counted overload ran past the body");
  }
}

#if USE_SODIUM {
  sodium_tests :: () {
    group("sodium_memzero wrapper");

    g: Guarded;

    prepare(*g, 0xcc);
    secure_memzero(body_of(*g), BODY);
    expect(all_bytes_are(body_of(*g), BODY, 0), "secure_memzero did not zero the body");
    expect(guards_intact(*g), "secure_memzero ran past the body");

    prepare(*g, 0xcc);
    secure_memzero(body_of(*g), 0);
    expect(all_bytes_are(body_of(*g), BODY, 0xcc), "secure_memzero with a zero count wrote something");

    prepare(*g, 0xcc);
    secure_memzero(body_of(*g), -1);
    expect(all_bytes_are(body_of(*g), BODY, 0xcc), "secure_memzero with a negative count wrote something");
  }
}

report :: () -> bool {
  emit("\n-- results by group --\n");

  for 0..group_count-1 {
    if group_fails[it] {
      emit("  % -- % checks, % FAILED\n", group_names[it], group_checks[it], group_fails[it]);
    } else {
      emit("  % -- % checks\n", group_names[it], group_checks[it]);
    }
  }

  emit("\n");
  if checks_failed == 0 {
    emit("ALL SECURE MEMORY TESTS PASSED  (% checks over % groups)\n", checks_run, group_count);
    return true;
  }
  emit("SECURE MEMORY TESTS FAILED  (% of % checks over % groups)\n", checks_failed, checks_run, group_count);
  return false;
}

#scope_export

// builder null writes to stdout as it goes; pass one to buffer the whole report so a
// threaded run does not interleave with other suites.
run_secure_memory_tests :: (builder: *String_Builder = null) -> bool {
  output_builder = builder;

  emit("=== Secure_Memory test suite ===\n");
  emit("USE_SODIUM = %, CPU = %\n", USE_SODIUM, CPU);

  memory_secure_overwrite_tests();
  secure_overwrite_tests();
  #if USE_SODIUM  sodium_tests();

  return report();
}