// SHA-256, HMAC-SHA256, PBKDF2-HMAC-SHA256 and OS randomness. Used for password hashing and session tokens.

Sha256 :: struct {
    state:  [8] u32;
    block:  [64] u8;
    used:   int;
    length: u64;
}

SHA256_K :: u32.[
    0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, 0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5,
    0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, 0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174,
    0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc, 0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da,
    0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7, 0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967,
    0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13, 0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85,
    0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3, 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070,
    0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5, 0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3,
    0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208, 0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2,
];

sha256_init :: (s: *Sha256) {
    s.state = .[0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, 0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19];
    s.used = 0;
    s.length = 0;
}

rotr :: inline (x: u32, n: u32) -> u32 {
    return (x >> n) | (x << (32 - n));
}

sha256_compress :: (s: *Sha256, block: *u8) {
    w: [64] u32 = ---;
    for i: 0..15 {
        p := block + i * 4;
        w[i] = ((cast(u32) p[0]) << 24) | ((cast(u32) p[1]) << 16) | ((cast(u32) p[2]) << 8) | (cast(u32) p[3]);
    }
    for i: 16..63 {
        s0 := rotr(w[i-15], 7) ^ rotr(w[i-15], 18) ^ (w[i-15] >> 3);
        s1 := rotr(w[i-2], 17) ^ rotr(w[i-2], 19)  ^ (w[i-2] >> 10);
        w[i] = w[i-16] + s0 + w[i-7] + s1;
    }
    a, b, c, d := s.state[0], s.state[1], s.state[2], s.state[3];
    e, f, g, h := s.state[4], s.state[5], s.state[6], s.state[7];
    for i: 0..63 {
        S1 := rotr(e, 6) ^ rotr(e, 11) ^ rotr(e, 25);
        ch := (e & f) ^ (~e & g);
        t1 := h + S1 + ch + SHA256_K[i] + w[i];
        S0 := rotr(a, 2) ^ rotr(a, 13) ^ rotr(a, 22);
        mj := (a & b) ^ (a & c) ^ (b & c);
        t2 := S0 + mj;
        h = g; g = f; f = e; e = d + t1;
        d = c; c = b; b = a; a = t1 + t2;
    }
    s.state[0] += a; s.state[1] += b; s.state[2] += c; s.state[3] += d;
    s.state[4] += e; s.state[5] += f; s.state[6] += g; s.state[7] += h;
}

sha256_update :: (s: *Sha256, data: *u8, count: int) {
    s.length += cast(u64) count;
    i := 0;
    while i < count {
        n := min(64 - s.used, count - i);
        memcpy(s.block.data + s.used, data + i, n);
        s.used += n;
        i += n;
        if s.used == 64 {
            sha256_compress(s, s.block.data);
            s.used = 0;
        }
    }
}

sha256_update :: inline (s: *Sha256, data: string) {
    sha256_update(s, data.data, data.count);
}

sha256_final :: (s: *Sha256) -> [32] u8 {
    bits := s.length * 8;
    pad: [72] u8;
    pad[0] = 0x80;
    pad_len := ifx s.used < 56 then 56 - s.used else 120 - s.used;
    for i: 0..7 pad[pad_len + i] = cast(u8) (bits >> cast(u64) ((7 - i) * 8));
    sha256_update(s, pad.data, pad_len + 8);
    result: [32] u8;
    for i: 0..7 {
        result[i*4 + 0] = cast(u8) (s.state[i] >> 24);
        result[i*4 + 1] = cast(u8) (s.state[i] >> 16);
        result[i*4 + 2] = cast(u8) (s.state[i] >> 8);
        result[i*4 + 3] = cast(u8)  s.state[i];
    }
    return result;
}

sha256 :: (data: string) -> [32] u8 {
    s: Sha256;
    sha256_init(*s);
    sha256_update(*s, data);
    return sha256_final(*s);
}

// HMAC with the inner/outer pad states precomputed, so PBKDF2 costs two compressions per iteration.
Hmac_Sha256 :: struct {
    inner, outer: Sha256;
}

hmac_init :: (h: *Hmac_Sha256, key: string) {
    k: [64] u8;
    if key.count > 64 {
        digest := sha256(key);
        memcpy(k.data, digest.data, 32);
    } else {
        memcpy(k.data, key.data, key.count);
    }
    ipad, opad: [64] u8;
    for i: 0..63 {
        ipad[i] = k[i] ^ 0x36;
        opad[i] = k[i] ^ 0x5c;
    }
    sha256_init(*h.inner);
    sha256_update(*h.inner, ipad.data, 64);
    sha256_init(*h.outer);
    sha256_update(*h.outer, opad.data, 64);
}

hmac_compute :: (h: *Hmac_Sha256, message: *u8, count: int) -> [32] u8 {
    inner := h.inner;
    sha256_update(*inner, message, count);
    digest := sha256_final(*inner);
    outer := h.outer;
    sha256_update(*outer, digest.data, 32);
    return sha256_final(*outer);
}

hmac_sha256 :: (key: string, message: string) -> [32] u8 {
    h: Hmac_Sha256;
    hmac_init(*h, key);
    return hmac_compute(*h, message.data, message.count);
}

// PBKDF2-HMAC-SHA256 with a 32-byte output (one block).
pbkdf2_sha256 :: (password: string, salt: string, iterations: int) -> [32] u8 {
    h: Hmac_Sha256;
    hmac_init(*h, password);
    first := alloc_string(salt.count + 4,, temp);
    memcpy(first.data, salt.data, salt.count);
    first[salt.count + 0] = 0;
    first[salt.count + 1] = 0;
    first[salt.count + 2] = 0;
    first[salt.count + 3] = 1;
    u := hmac_compute(*h, first.data, first.count);
    result := u;
    for 2..iterations {
        u = hmac_compute(*h, u.data, 32);
        for i: 0..31 result[i] ^= u[i];
    }
    return result;
}

constant_time_equal :: (a: string, b: string) -> bool {
    if a.count != b.count return false;
    diff: u8;
    for i: 0..a.count-1 diff |= a[i] ^ b[i];
    return diff == 0;
}

to_hex :: (data: *u8, count: int) -> string {
    DIGITS :: "0123456789abcdef";
    result := alloc_string(count * 2);
    for i: 0..count-1 {
        result[i*2]     = DIGITS[data[i] >> 4];
        result[i*2 + 1] = DIGITS[data[i] & 0xF];
    }
    return result;
}

as_string :: (bytes: *[32] u8) -> string {
    s: string;
    s.data  = bytes.data;
    s.count = 32;
    return s;
}

//
// Randomness from the OS.
//

#if OS == .WINDOWS {
    tf_bcrypt :: #system_library "bcrypt";
    BCryptGenRandom :: (algorithm: *void, buffer: *u8, count: u32, flags: u32) -> s32 #foreign tf_bcrypt;

    random_bytes :: (buffer: *u8, count: int) -> bool {
        BCRYPT_USE_SYSTEM_PREFERRED_RNG :: 2;
        return BCryptGenRandom(null, buffer, cast(u32) count, BCRYPT_USE_SYSTEM_PREFERRED_RNG) == 0;
    }
} else {
    tf_libc :: #system_library "libc";
    getrandom :: (buffer: *void, count: u64, flags: u32) -> s64 #foreign tf_libc;

    random_bytes :: (buffer: *u8, count: int) -> bool {
        done := 0;
        while done < count {
            n := getrandom(buffer + done, cast(u64) (count - done), 0);
            if n <= 0 return false;
            done += n;
        }
        return true;
    }
}