// Copyright (C) 2026 Douglas Quigg (dstroy0) <dquigg123@gmail.com>
// SPDX-License-Identifier: AGPL-3.0-or-later
//
// Unit tests for the RFC 1951 INFLATE core (network_drivers/presentation/inflate).
// Vectors are generated with Python's zlib (raw deflate, wbits=-15) so the
// decompressor is checked against the reference implementation. The
// permessage-deflate vectors are SYNC_FLUSH'd with the trailing 00 00 ff ff
// stripped (RFC 7692 sender); the test appends the marker back before inflating,
// mirroring the receiver.

#include "network_drivers/presentation/inflate/inflate.h"
#include <string.h>
#include <unity.h>

// --- vectors (generated by scratchpad/gen_vectors.py against zlib) ----------
static const uint8_t k_hello_in[] = {243, 72, 205, 201, 201, 215, 81, 8, 207, 47, 202, 73, 81, 4, 0};
static const uint8_t k_hello_out[] = {72, 101, 108, 108, 111, 44, 32, 87, 111, 114, 108, 100, 33};

static const uint8_t k_repeat_in[] = {115, 116, 114, 118, 196, 141, 0};
static const uint8_t k_repeat_out[] = {65, 66, 67, 65, 66, 67, 65, 66, 67, 65, 66, 67, 65, 66, 67,
                                       65, 66, 67, 65, 66, 67, 65, 66, 67, 65, 66, 67, 65, 66, 67};

static const uint8_t k_stored_in[] = {1,  25,  0,   230, 255, 0,   1,   2,  120, 121, 122, 32,  98,  105, 110,
                                      97, 114, 121, 45,  105, 115, 104, 32, 112, 97,  121, 108, 111, 97,  100};
static const uint8_t k_stored_out[] = {0,  1,   2,   120, 121, 122, 32, 98,  105, 110, 97, 114, 121,
                                       45, 105, 115, 104, 32,  112, 97, 121, 108, 111, 97, 100};

static const uint8_t k_dyn_in[] = {45,  141, 219, 17,  195, 32,  12,  4,   91,  185, 52,  224, 158, 32,  22, 160,
                                   4,   35,  155, 167, 161, 250, 104, 60,  249, 222, 189, 189, 26,  8,   87, 227,
                                   247, 23,  54,  203, 72,  112, 114, 227, 211, 142, 179, 64,  58,  101, 84, 197,
                                   209, 172, 137, 93,  252, 134, 211, 168, 119, 76,  88,  149, 6,   215, 0,  199,
                                   157, 20,  45,  74,  136, 124, 53,  201, 186, 245, 101, 67,  144, 129, 78, 55,
                                   39,  31,  231, 63,  191, 27,  87,  177, 200, 102, 83,  158, 131, 215, 15};
static const uint8_t k_dyn_out[] = {
    116, 104, 101, 32,  113, 117, 105, 99,  107, 32,  98,  114, 111, 119, 110, 32,  102, 111, 120, 32,  106,
    117, 109, 112, 115, 32,  111, 118, 101, 114, 32,  116, 104, 101, 32,  108, 97,  122, 121, 32,  100, 111,
    103, 46,  32,  112, 97,  99,  107, 32,  109, 121, 32,  98,  111, 120, 32,  119, 105, 116, 104, 32,  102,
    105, 118, 101, 32,  100, 111, 122, 101, 110, 32,  108, 105, 113, 117, 111, 114, 32,  106, 117, 103, 115,
    46,  32,  104, 111, 119, 32,  118, 101, 120, 105, 110, 103, 108, 121, 32,  113, 117, 105, 99,  107, 32,
    100, 97,  102, 116, 32,  122, 101, 98,  114, 97,  115, 32,  106, 117, 109, 112, 33};

static const uint8_t k_empty_in[] = {3, 0};

static const uint8_t k_pmd_hello_in[] = {242, 72, 205, 201, 201, 215, 81, 8, 207, 47, 202, 73, 81, 4, 0};
static const uint8_t k_pmd_hello_out[] = {72, 101, 108, 108, 111, 44, 32, 87, 111, 114, 108, 100, 33};

static const uint8_t k_pmd_repeat_in[] = {114, 116, 114, 118, 196, 64, 0, 0};
static const uint8_t k_pmd_repeat_out[] = {65, 66, 67, 65, 66, 67, 65, 66, 67, 65, 66,
                                           67, 65, 66, 67, 65, 66, 67, 65, 66, 67};

// Malformed DEFLATE bit-streams (scratchpad/gen_inflate_bad.py): each drives a distinct
// error guard in the fixed/dynamic block decoders. All must be rejected as MALFORMED.
static const uint8_t k_bad_lencode_286[] = {115, 28, 3};             // fixed block emits length symbol 286
static const uint8_t k_bad_distcode_30[] = {115, 4, 62};             // distance symbol 30 (>= 30)
static const uint8_t k_bad_dist_gt_out[] = {115, 4, 66};             // back-ref distance past start of output
static const uint8_t k_trunc_len_extra[] = {115, 68, 2};             // length code needs extra bits, stream ends
static const uint8_t k_trunc_dist_extra[] = {115, 4, 18};            // distance code needs extra bits, stream ends
static const uint8_t k_stored_len_over[] = {1, 10, 0, 245, 255, 65}; // stored LEN exceeds the input
static const uint8_t k_dyn_hdr_trunc[] = {5};                        // dynamic header truncated mid HLIT/HDIST/HCLEN
static const uint8_t k_dyn_hlit_big[] = {253, 0, 0};                 // oversized dynamic header

// Dynamic-block code-length / table-construction guards (scratchpad/gen_inflate_dyn.py).
static const uint8_t k_dyn_repeat16_at0[] = {5, 224, 3, 4, 0, 0, 0, 0, 0, 4}; // repeat-16 with no previous length
static const uint8_t k_dyn_repeat_past_end[] = {5, 224, 129, 4,   0,   0,
                                                0, 0,   0,   252, 255, 3};          // repeat runs past the code count
static const uint8_t k_dyn_cl_incomplete[] = {5, 224, 1, 4, 0, 0, 0, 0, 0, 0};      // code-length code is incomplete
static const uint8_t k_dyn_cl_oversub[] = {5, 224, 1, 4, 0, 0, 0, 64, 16, 0};       // code-length code over-subscribed
static const uint8_t k_dyn_no_eob[] = {5, 224, 129, 4, 0, 0, 0, 0, 0, 252, 111, 3}; // no end-of-block code
static const uint8_t k_dyn_cl_decode_trunc[] = {5, 224, 21, 8, 0,
                                                0, 0,   0,  0, 0}; // truncated before any code-length symbol
static const uint8_t k_dyn_empty_dist[] = {253, 224, 129, 0,   0,   0,   0,
                                           0,   16,  180, 249, 159, 178, 2}; // literals-only: empty distance code
static const uint8_t k_dyn_litlen_incomplete[] = {253, 224, 1,   1,   0,  0,  0, 128,
                                                  16,  108, 235, 255, 41, 87, 0}; // incomplete literal/length code
static const uint8_t k_dyn_dist_incomplete[] = {253, 225, 1,   9,   0,  0,   0, 128,
                                                32,  108, 243, 255, 41, 167, 4};    // incomplete distance code
static const uint8_t k_len_extra_aligned[] = {59, 113, 226, 196, 137, 19, 39, 144}; // length extra bits past end
static const uint8_t k_dist_extra_aligned[] = {59, 1, 36};                          // distance extra bits past end
static const uint8_t k_header_leftover_bits[] = {58, 113, 226, 196, 137, 19, 0};    // next block header needs a byte
static const uint8_t k_dyn_decode_hole[] = {5, 224, 1, 5, 0, 0, 0, 0, 32, 252, 127, 157, 255, 63}; // code past MAXBITS
static const uint8_t k_dyn_literal_len[] = {5, 224, 1, 4, 0, 0, 0, 0, 16, 4}; // literal code length then truncated

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

static uint8_t g_scratch[INFLATE_SCRATCH_SIZE];
static uint8_t g_out[1024];

void setUp()
{
    memset(g_out, 0, sizeof(g_out));
}
void tearDown()
{
}

// Inflate src and assert it equals expected.
static void expect_inflates_to(const uint8_t *src, size_t src_len, const uint8_t *exp, size_t exp_len)
{
    size_t out_len = 0;
    int rc = inflate_raw(src, src_len, g_out, sizeof(g_out), &out_len, g_scratch, sizeof(g_scratch));
    TEST_ASSERT_EQUAL_INT(INFLATE_OK, rc);
    TEST_ASSERT_EQUAL_size_t(exp_len, out_len);
    if (exp_len)
        TEST_ASSERT_EQUAL_MEMORY(exp, g_out, exp_len);
}

void test_fixed_huffman()
{
    expect_inflates_to(k_hello_in, sizeof(k_hello_in), k_hello_out, sizeof(k_hello_out));
}

void test_back_references()
{
    expect_inflates_to(k_repeat_in, sizeof(k_repeat_in), k_repeat_out, sizeof(k_repeat_out));
}

void test_stored_block()
{
    expect_inflates_to(k_stored_in, sizeof(k_stored_in), k_stored_out, sizeof(k_stored_out));
}

void test_dynamic_huffman()
{
    expect_inflates_to(k_dyn_in, sizeof(k_dyn_in), k_dyn_out, sizeof(k_dyn_out));
}

void test_empty_message()
{
    expect_inflates_to(k_empty_in, sizeof(k_empty_in), nullptr, 0);
}

// permessage-deflate: append the 00 00 ff ff marker, then inflate (RFC 7692).
static void expect_pmd_inflates_to(const uint8_t *src, size_t src_len, const uint8_t *exp, size_t exp_len)
{
    uint8_t buf[256];
    TEST_ASSERT_TRUE(src_len + 4 <= sizeof(buf));
    memcpy(buf, src, src_len);
    buf[src_len] = 0x00;
    buf[src_len + 1] = 0x00;
    buf[src_len + 2] = 0xff;
    buf[src_len + 3] = 0xff;
    expect_inflates_to(buf, src_len + 4, exp, exp_len);
}

void test_permessage_deflate_marker()
{
    expect_pmd_inflates_to(k_pmd_hello_in, sizeof(k_pmd_hello_in), k_pmd_hello_out, sizeof(k_pmd_hello_out));
}

void test_permessage_deflate_back_references()
{
    expect_pmd_inflates_to(k_pmd_repeat_in, sizeof(k_pmd_repeat_in), k_pmd_repeat_out, sizeof(k_pmd_repeat_out));
}

void test_output_overflow_fails_closed()
{
    size_t out_len = 0;
    uint8_t tiny[5];
    int rc = inflate_raw(k_hello_in, sizeof(k_hello_in), tiny, sizeof(tiny), &out_len, g_scratch, sizeof(g_scratch));
    TEST_ASSERT_EQUAL_INT(INFLATE_ERR_OVERFLOW, rc);
}

void test_scratch_too_small_fails_closed()
{
    size_t out_len = 0;
    uint8_t small[INFLATE_SCRATCH_SIZE - 1];
    int rc = inflate_raw(k_hello_in, sizeof(k_hello_in), g_out, sizeof(g_out), &out_len, small, sizeof(small));
    TEST_ASSERT_EQUAL_INT(INFLATE_ERR_SCRATCH, rc);
}

void test_truncated_input_is_malformed()
{
    size_t out_len = 0;
    // Half of the fixed-Huffman stream: decode runs out of input mid-symbol.
    int rc =
        inflate_raw(k_hello_in, sizeof(k_hello_in) / 2, g_out, sizeof(g_out), &out_len, g_scratch, sizeof(g_scratch));
    TEST_ASSERT_EQUAL_INT(INFLATE_ERR_MALFORMED, rc);
}

void test_reserved_block_type_is_malformed()
{
    size_t out_len = 0;
    const uint8_t bad[] = {0x06}; // BFINAL=0, BTYPE=11 (reserved)
    int rc = inflate_raw(bad, sizeof(bad), g_out, sizeof(g_out), &out_len, g_scratch, sizeof(g_scratch));
    TEST_ASSERT_EQUAL_INT(INFLATE_ERR_MALFORMED, rc);
}

void test_corrupt_stored_nlen_is_malformed()
{
    uint8_t bad[sizeof(k_stored_in)];
    memcpy(bad, k_stored_in, sizeof(bad));
    bad[3] ^= 0xFF; // break the NLEN ones-complement check
    size_t out_len = 0;
    int rc = inflate_raw(bad, sizeof(bad), g_out, sizeof(g_out), &out_len, g_scratch, sizeof(g_scratch));
    TEST_ASSERT_EQUAL_INT(INFLATE_ERR_MALFORMED, rc);
}

void test_inflate_error_paths()
{
    uint8_t out[64];
    size_t out_len = 0;
    // OVERFLOW: a valid stream decompressed into a buffer that is too small.
    TEST_ASSERT_EQUAL_INT(INFLATE_ERR_OVERFLOW,
                          inflate_raw(k_hello_in, sizeof(k_hello_in), out, 4, &out_len, g_scratch, sizeof(g_scratch)));
    // A stored block into a too-small buffer also overflows.
    TEST_ASSERT_EQUAL_INT(INFLATE_ERR_OVERFLOW, inflate_raw(k_stored_in, sizeof(k_stored_in), out, 3, &out_len,
                                                            g_scratch, sizeof(g_scratch)));
    // Reserved block type (BTYPE = 11) is malformed.
    const uint8_t bad_btype[2] = {0x06, 0x00};
    TEST_ASSERT_EQUAL_INT(INFLATE_ERR_MALFORMED, inflate_raw(bad_btype, sizeof(bad_btype), out, sizeof(out), &out_len,
                                                             g_scratch, sizeof(g_scratch)));
    // Stored block whose NLEN is not the ones-complement of LEN.
    const uint8_t bad_stored[5] = {0x01, 0x03, 0x00, 0x00, 0x00}; // BFINAL=1 BTYPE=00 LEN=3 NLEN=0
    TEST_ASSERT_EQUAL_INT(INFLATE_ERR_MALFORMED, inflate_raw(bad_stored, sizeof(bad_stored), out, sizeof(out), &out_len,
                                                             g_scratch, sizeof(g_scratch)));
    // A stored block truncated before its LEN/NLEN.
    const uint8_t trunc_stored[1] = {0x00};
    TEST_ASSERT_TRUE(inflate_raw(trunc_stored, 1, out, sizeof(out), &out_len, g_scratch, sizeof(g_scratch)) !=
                     INFLATE_OK);
    // A valid compressed stream cut short runs out of input mid-decode.
    TEST_ASSERT_TRUE(inflate_raw(k_hello_in, 3, out, sizeof(out), &out_len, g_scratch, sizeof(g_scratch)) !=
                     INFLATE_OK);
}

// Each crafted malformed stream is rejected by the specific decoder guard it targets.
void test_malformed_deflate_blocks()
{
    uint8_t out[64];
    size_t out_len = 0;
#define BAD(v)                                                                                                         \
    TEST_ASSERT_EQUAL_INT(INFLATE_ERR_MALFORMED,                                                                       \
                          inflate_raw(v, sizeof(v), out, sizeof(out), &out_len, g_scratch, sizeof(g_scratch)))
    BAD(k_bad_lencode_286);
    BAD(k_bad_distcode_30);
    BAD(k_bad_dist_gt_out);
    BAD(k_trunc_len_extra);
    BAD(k_trunc_dist_extra);
    BAD(k_stored_len_over);
    BAD(k_dyn_hdr_trunc);
    BAD(k_dyn_hlit_big);
    BAD(k_dyn_repeat16_at0);
    BAD(k_dyn_repeat_past_end);
    BAD(k_dyn_cl_incomplete);
    BAD(k_dyn_cl_oversub);
    BAD(k_dyn_no_eob);
    BAD(k_dyn_cl_decode_trunc);
    BAD(k_dyn_empty_dist); // 120 (empty dist code) hit during construction; the truncated data then fails
    BAD(k_dyn_litlen_incomplete);
    BAD(k_dyn_dist_incomplete);
    BAD(k_len_extra_aligned);
    BAD(k_dist_extra_aligned);
    BAD(k_header_leftover_bits);
    BAD(k_dyn_decode_hole);
    BAD(k_dyn_literal_len);
#undef BAD
    // A back-reference whose copy length overflows the remaining output buffer -> OVERFLOW.
    TEST_ASSERT_EQUAL_INT(INFLATE_ERR_OVERFLOW, inflate_raw(k_repeat_in, sizeof(k_repeat_in), out, 4, &out_len,
                                                            g_scratch, sizeof(g_scratch)));
}

int main()
{
    UNITY_BEGIN();
    RUN_TEST(test_malformed_deflate_blocks);
    RUN_TEST(test_fixed_huffman);
    RUN_TEST(test_back_references);
    RUN_TEST(test_stored_block);
    RUN_TEST(test_dynamic_huffman);
    RUN_TEST(test_empty_message);
    RUN_TEST(test_permessage_deflate_marker);
    RUN_TEST(test_permessage_deflate_back_references);
    RUN_TEST(test_output_overflow_fails_closed);
    RUN_TEST(test_scratch_too_small_fails_closed);
    RUN_TEST(test_truncated_input_is_malformed);
    RUN_TEST(test_reserved_block_type_is_malformed);
    RUN_TEST(test_corrupt_stored_nlen_is_malformed);
    RUN_TEST(test_inflate_error_paths);
    return UNITY_END();
}
