/*
 * Unit tests for validity TLS session management functions
 *
 * Copyright (C) 2024 libfprint contributors
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 */

#include <glib.h>
#include <string.h>

#include <openssl/evp.h>
#include <openssl/hmac.h>
#include <openssl/rand.h>

#include "fpi-device.h"
#include "fpi-ssm.h"
#include "fpi-byte-reader.h"

/* We include the TLS header and use function declarations directly.
 * The test links against the driver static lib. */
#include "drivers/validity/validity_tls.h"
#include "drivers/validity/vcsfw_protocol.h"

/* ================================================================
 * Test: PRF produces deterministic output
 * ================================================================ */
static void
test_prf_deterministic (void)
{
  guint8 secret[] = { 0x01, 0x02, 0x03, 0x04 };
  guint8 seed[] = { 0x05, 0x06, 0x07, 0x08 };
  guint8 output1[48];
  guint8 output2[48];

  validity_tls_prf (secret, sizeof (secret), seed, sizeof (seed),
                    output1, sizeof (output1));
  validity_tls_prf (secret, sizeof (secret), seed, sizeof (seed),
                    output2, sizeof (output2));

  g_assert_cmpmem (output1, sizeof (output1), output2, sizeof (output2));
}

/* ================================================================
 * Test: PRF with known TLS 1.2 test vector
 * ================================================================
 * RFC 5246 does not define test vectors for SHA-256 PRF directly,
 * but we verify our implementation against python-validity's output.
 */
static void
test_prf_output_length (void)
{
  guint8 secret[32];
  guint8 seed[64];
  guint8 output[0x120]; /* Same as key_block size */

  memset (secret, 0xAB, sizeof (secret));
  memset (seed, 0xCD, sizeof (seed));

  validity_tls_prf (secret, sizeof (secret), seed, sizeof (seed),
                    output, sizeof (output));

  /* PRF output should not be all zeros */
  gboolean all_zero = TRUE;
  for (gsize i = 0; i < sizeof (output); i++)
    {
      if (output[i] != 0)
        {
          all_zero = FALSE;
          break;
        }
    }
  g_assert_false (all_zero);
}

/* ================================================================
 * Test: PRF with different lengths uses correct number of HMAC iters
 * ================================================================ */
static void
test_prf_short_output (void)
{
  guint8 secret[] = { 0x01 };
  guint8 seed[] = { 0x02 };
  guint8 output_short[16];
  guint8 output_long[48];

  validity_tls_prf (secret, 1, seed, 1, output_short, sizeof (output_short));
  validity_tls_prf (secret, 1, seed, 1, output_long, sizeof (output_long));

  /* First 16 bytes should match */
  g_assert_cmpmem (output_short, 16, output_long, 16);
}

/* ================================================================
 * Test: Encrypt then decrypt roundtrip
 * ================================================================ */
static void
test_encrypt_decrypt_roundtrip (void)
{
  ValidityTlsState tls;
  validity_tls_init (&tls);

  /* Set up encryption/decryption keys (same for roundtrip test) */
  memset (tls.encryption_key, 0x42, TLS_AES_KEY_SIZE);
  memset (tls.decryption_key, 0x42, TLS_AES_KEY_SIZE);

  guint8 plaintext[] = "Hello, TLS! This is a test message for encryption.";
  gsize pt_len = sizeof (plaintext);

  gsize enc_len;
  guint8 *encrypted = validity_tls_encrypt (&tls, plaintext, pt_len, &enc_len);
  g_assert_nonnull (encrypted);
  g_assert_cmpuint (enc_len, >, pt_len); /* IV + padded ciphertext */

  GError *error = NULL;
  gsize dec_len;
  guint8 *decrypted = validity_tls_decrypt (&tls, encrypted, enc_len,
                                            &dec_len, &error);
  g_assert_no_error (error);
  g_assert_nonnull (decrypted);
  g_assert_cmpmem (plaintext, pt_len, decrypted, dec_len);

  g_free (encrypted);
  g_free (decrypted);
  validity_tls_free (&tls);
}

/* ================================================================
 * Test: Encrypt with block-aligned data
 * ================================================================ */
static void
test_encrypt_block_aligned (void)
{
  ValidityTlsState tls;
  validity_tls_init (&tls);

  memset (tls.encryption_key, 0x55, TLS_AES_KEY_SIZE);
  memset (tls.decryption_key, 0x55, TLS_AES_KEY_SIZE);

  /* 16 bytes = exactly one AES block */
  guint8 plaintext[16];
  memset (plaintext, 0xAA, 16);

  gsize enc_len;
  guint8 *encrypted = validity_tls_encrypt (&tls, plaintext, 16, &enc_len);
  g_assert_nonnull (encrypted);
  /* Should be IV(16) + 32 bytes (16 data + 16 padding since pad=0x0f*16) */
  g_assert_cmpuint (enc_len, ==, 16 + 32);

  GError *error = NULL;
  gsize dec_len;
  guint8 *decrypted = validity_tls_decrypt (&tls, encrypted, enc_len,
                                            &dec_len, &error);
  g_assert_no_error (error);
  g_assert_nonnull (decrypted);
  g_assert_cmpuint (dec_len, ==, 16);
  g_assert_cmpmem (plaintext, 16, decrypted, 16);

  g_free (encrypted);
  g_free (decrypted);
  validity_tls_free (&tls);
}

/* ================================================================
 * Test: Decrypt with invalid data fails
 * ================================================================ */
static void
test_decrypt_invalid (void)
{
  ValidityTlsState tls;
  validity_tls_init (&tls);

  memset (tls.decryption_key, 0x55, TLS_AES_KEY_SIZE);

  /* Too short for IV + block */
  guint8 short_data[10];
  memset (short_data, 0, sizeof (short_data));

  GError *error = NULL;
  gsize dec_len;
  guint8 *decrypted = validity_tls_decrypt (&tls, short_data,
                                            sizeof (short_data),
                                            &dec_len, &error);
  g_assert_null (decrypted);
  g_assert_error (error, FP_DEVICE_ERROR, FP_DEVICE_ERROR_PROTO);
  g_clear_error (&error);

  validity_tls_free (&tls);
}

/* ================================================================
 * Test: PSK derivation runs without crashing
 * ================================================================ */
static void
test_psk_derivation (void)
{
  ValidityTlsState tls;
  validity_tls_init (&tls);

  validity_tls_derive_psk (&tls);

  /* PSK keys should not be all zeros */
  gboolean all_zero = TRUE;
  for (gsize i = 0; i < TLS_AES_KEY_SIZE; i++)
    {
      if (tls.psk_encryption_key[i] != 0)
        {
          all_zero = FALSE;
          break;
        }
    }
  g_assert_false (all_zero);

  all_zero = TRUE;
  for (gsize i = 0; i < TLS_AES_KEY_SIZE; i++)
    {
      if (tls.psk_validation_key[i] != 0)
        {
          all_zero = FALSE;
          break;
        }
    }
  g_assert_false (all_zero);

  validity_tls_free (&tls);
}

/* ================================================================
 * Test: PSK derivation is deterministic
 * ================================================================ */
static void
test_psk_deterministic (void)
{
  ValidityTlsState tls1, tls2;
  validity_tls_init (&tls1);
  validity_tls_init (&tls2);

  validity_tls_derive_psk (&tls1);
  validity_tls_derive_psk (&tls2);

  g_assert_cmpmem (tls1.psk_encryption_key, TLS_AES_KEY_SIZE,
                   tls2.psk_encryption_key, TLS_AES_KEY_SIZE);
  g_assert_cmpmem (tls1.psk_validation_key, TLS_AES_KEY_SIZE,
                   tls2.psk_validation_key, TLS_AES_KEY_SIZE);

  validity_tls_free (&tls1);
  validity_tls_free (&tls2);
}

/* ================================================================
 * Test: Flash parse with empty data fails gracefully
 * ================================================================ */
static void
test_flash_parse_empty (void)
{
  ValidityTlsState tls;
  validity_tls_init (&tls);

  GError *error = NULL;
  guint8 empty_flash[] = { 0xFF, 0xFF, 0x00, 0x00 }; /* end block */

  /* Flash with only end marker → missing keys */
  gboolean result = validity_tls_parse_flash (&tls, empty_flash,
                                              sizeof (empty_flash),
                                              &error);
  g_assert_false (result);
  g_assert_nonnull (error);
  g_assert_false (tls.keys_loaded);

  g_clear_error (&error);
  validity_tls_free (&tls);
}

/* ================================================================
 * Test: Flash parse with truncated data fails gracefully
 * ================================================================ */
static void
test_flash_parse_truncated (void)
{
  ValidityTlsState tls;
  validity_tls_init (&tls);

  GError *error = NULL;
  guint8 truncated[] = { 0x03, 0x00, 0xFF, 0x00 }; /* cert block w/ impossibly large size */

  gboolean result = validity_tls_parse_flash (&tls, truncated,
                                              sizeof (truncated),
                                              &error);
  /* Should fail due to block size exceeding remaining data */
  g_assert_false (result);
  g_assert_nonnull (error);
  g_clear_error (&error);

  validity_tls_free (&tls);
}

/* ================================================================
 * Test: Init/free cycle doesn't leak
 * ================================================================ */
static void
test_init_free (void)
{
  ValidityTlsState tls;

  for (int i = 0; i < 10; i++)
    {
      validity_tls_init (&tls);
      validity_tls_free (&tls);
    }
}

/* ================================================================
 * Test: Build ClientHello produces valid TLS record
 * ================================================================ */
static void
test_build_client_hello (void)
{
  ValidityTlsState tls;
  validity_tls_init (&tls);

  gsize out_len;
  guint8 *hello = validity_tls_build_client_hello (&tls, &out_len);

  g_assert_nonnull (hello);
  g_assert_cmpuint (out_len, >, 4 + 5); /* prefix(4) + record header(5) minimum */

  /* Check prefix: 0x44 0x00 0x00 0x00 */
  g_assert_cmpint (hello[0], ==, 0x44);
  g_assert_cmpint (hello[1], ==, 0x00);
  g_assert_cmpint (hello[2], ==, 0x00);
  g_assert_cmpint (hello[3], ==, 0x00);

  /* Check TLS record header */
  g_assert_cmpint (hello[4], ==, 0x16); /* handshake */
  g_assert_cmpint (hello[5], ==, 0x03); /* version major */
  g_assert_cmpint (hello[6], ==, 0x03); /* version minor */

  /* client_random should have been set */
  gboolean has_random = FALSE;
  for (gsize i = 0; i < TLS_RANDOM_SIZE; i++)
    {
      if (tls.client_random[i] != 0)
        {
          has_random = TRUE;
          break;
        }
    }
  g_assert_true (has_random);

  g_free (hello);
  validity_tls_free (&tls);
}

/* ================================================================
 * Test: Wrap/unwrap with invalid data fails gracefully
 * ================================================================ */
static void
test_unwrap_invalid (void)
{
  ValidityTlsState tls;
  validity_tls_init (&tls);

  GError *error = NULL;
  gsize out_len;

  /* Short data → truncated record header */
  guint8 short_data[] = { 0x17, 0x03 };
  guint8 *result = validity_tls_unwrap_response (&tls, short_data,
                                                  sizeof (short_data),
                                                  &out_len, &error);
  g_assert_null (result);
  g_assert_nonnull (error);
  g_clear_error (&error);

  /* App data before secure channel */
  guint8 app_early[] = { 0x17, 0x03, 0x03, 0x00, 0x10,
                         0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
                         0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f };
  result = validity_tls_unwrap_response (&tls, app_early,
                                          sizeof (app_early),
                                          &out_len, &error);
  g_assert_null (result);
  g_assert_nonnull (error);
  g_clear_error (&error);

  validity_tls_free (&tls);
}

/* ================================================================
 * Regression: Bug #1 — Flash parse requires PSK for private key
 *
 * Private key block (ID 4) is encrypted with PSK. Calling parse_flash
 * without first deriving PSK must fail (HMAC mismatch), proving the
 * ordering dependency. This catches the bug where flash_read SSM
 * parsed flash data BEFORE PSK derivation had occurred.
 * ================================================================ */
static void
test_flash_parse_needs_psk (void)
{
  ValidityTlsState tls_with_psk, tls_no_psk;
  validity_tls_init (&tls_with_psk);
  validity_tls_init (&tls_no_psk);

  /* Derive PSK so we can build a valid encrypted private key block */
  validity_tls_derive_psk (&tls_with_psk);

  /* Build a realistic flash image with a cert block + encrypted privkey block.
   * We use a minimal cert (just 16 bytes of dummy data) and a privkey block
   * that's encrypted with the proper PSK. */

  /* Step 1: Build a cert body */
  guint8 cert_body[16];
  memset (cert_body, 0xAA, sizeof (cert_body));

  /* Step 2: Build a private-key body encrypted with PSK */
  guint8 priv_plaintext[96]; /* d(32) + pad for block alignment */
  memset (priv_plaintext, 0xBB, sizeof (priv_plaintext));

  /* Encrypt plaintext with PSK encryption key */
  guint8 iv[TLS_IV_SIZE];
  memset (iv, 0x11, TLS_IV_SIZE);
  gsize ct_len = sizeof (priv_plaintext);
  guint8 *ciphertext = g_malloc (TLS_IV_SIZE + ct_len);
  memcpy (ciphertext, iv, TLS_IV_SIZE);

  EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new ();
  int out_len, final_len;
  EVP_EncryptInit_ex (ctx, EVP_aes_256_cbc (), NULL,
                      tls_with_psk.psk_encryption_key, iv);
  EVP_CIPHER_CTX_set_padding (ctx, 0);
  EVP_EncryptUpdate (ctx, ciphertext + TLS_IV_SIZE, &out_len,
                     priv_plaintext, ct_len);
  EVP_EncryptFinal_ex (ctx, ciphertext + TLS_IV_SIZE + out_len, &final_len);
  EVP_CIPHER_CTX_free (ctx);
  gsize enc_total = TLS_IV_SIZE + ct_len;

  /* HMAC over (iv + ciphertext) with psk_validation_key */
  guint8 mac[TLS_HMAC_SIZE];
  unsigned int mac_len;
  HMAC (EVP_sha256 (),
        tls_with_psk.psk_validation_key, TLS_AES_KEY_SIZE,
        ciphertext, enc_total, mac, &mac_len);

  /* Private key block payload: 0x02 || ciphertext || hmac */
  gsize priv_block_len = 1 + enc_total + TLS_HMAC_SIZE;
  guint8 *priv_block = g_malloc (priv_block_len);
  priv_block[0] = 0x02;
  memcpy (priv_block + 1, ciphertext, enc_total);
  memcpy (priv_block + 1 + enc_total, mac, TLS_HMAC_SIZE);
  g_free (ciphertext);

  /* Build flash image: [cert_header][cert_body][priv_header][priv_body][end] */
  GByteArray *flash = g_byte_array_new ();

  /* Cert block header: id=0x0003, size, sha256 hash */
  guint8 cert_hdr[TLS_FLASH_BLOCK_HEADER_SIZE];
  FP_WRITE_UINT16_LE (cert_hdr, TLS_FLASH_BLOCK_CERT);
  FP_WRITE_UINT16_LE (cert_hdr + 2, sizeof (cert_body));
  GChecksum *cs = g_checksum_new (G_CHECKSUM_SHA256);
  gsize hash_len = 32;
  g_checksum_update (cs, cert_body, sizeof (cert_body));
  g_checksum_get_digest (cs, cert_hdr + 4, &hash_len);
  g_checksum_free (cs);
  g_byte_array_append (flash, cert_hdr, sizeof (cert_hdr));
  g_byte_array_append (flash, cert_body, sizeof (cert_body));

  /* Priv block header */
  guint8 priv_hdr[TLS_FLASH_BLOCK_HEADER_SIZE];
  FP_WRITE_UINT16_LE (priv_hdr, TLS_FLASH_BLOCK_PRIVKEY);
  FP_WRITE_UINT16_LE (priv_hdr + 2, priv_block_len);
  cs = g_checksum_new (G_CHECKSUM_SHA256);
  hash_len = 32;
  g_checksum_update (cs, priv_block, priv_block_len);
  g_checksum_get_digest (cs, priv_hdr + 4, &hash_len);
  g_checksum_free (cs);
  g_byte_array_append (flash, priv_hdr, sizeof (priv_hdr));
  g_byte_array_append (flash, priv_block, priv_block_len);

  /* End marker */
  guint8 end_marker[4] = { 0xFF, 0xFF, 0x00, 0x00 };
  g_byte_array_append (flash, end_marker, sizeof (end_marker));

  /* TEST: Without PSK, parse_flash must fail on the privkey block */
  GError *error = NULL;
  gboolean result = validity_tls_parse_flash (&tls_no_psk,
                                              flash->data, flash->len,
                                              &error);
  g_assert_false (result);
  g_assert_nonnull (error);
  /* Should fail with HMAC-related error since PSK is all zeros */
  g_clear_error (&error);

  g_byte_array_free (flash, TRUE);
  g_free (priv_block);
  validity_tls_free (&tls_with_psk);
  validity_tls_free (&tls_no_psk);
}

/* ================================================================
 * Regression: Bug #2 — READ_FLASH command format
 *
 * The READ_FLASH command must be exactly 13 bytes matching
 * python-validity: pack('<BBBHLL', 0x40, partition, 1, 0, addr, size).
 * Old code only sent 10 bytes, missing the access flag and reserved
 * field. This test verifies the command constant and expected layout.
 * ================================================================ */
static void
test_flash_cmd_format (void)
{
  /* Verify the command byte */
  g_assert_cmpint (VCSFW_CMD_READ_FLASH, ==, 0x40);

  /* Build the same command as validity_tls_flash_read_run_state does */
  guint8 cmd[13];
  cmd[0] = VCSFW_CMD_READ_FLASH;
  cmd[1] = 0x01;           /* partition */
  cmd[2] = 0x01;           /* access flag */
  FP_WRITE_UINT16_LE (&cmd[3], 0x0000);  /* reserved */
  FP_WRITE_UINT32_LE (&cmd[5], 0x0000);  /* offset */
  FP_WRITE_UINT32_LE (&cmd[9], 0x1000);  /* size */

  /* Verify total size is 13 (not 10 like the old bug) */
  g_assert_cmpuint (sizeof (cmd), ==, 13);

  /* Verify byte layout matches python-validity's pack('<BBBHLL', ...) */
  g_assert_cmpint (cmd[0], ==, 0x40);  /* command */
  g_assert_cmpint (cmd[1], ==, 0x01);  /* partition */
  g_assert_cmpint (cmd[2], ==, 0x01);  /* access flag (was missing) */
  g_assert_cmpint (cmd[3], ==, 0x00);  /* reserved lo */
  g_assert_cmpint (cmd[4], ==, 0x00);  /* reserved hi */
  /* offset at bytes 5-8 (LE uint32 = 0) */
  g_assert_cmpint (cmd[5], ==, 0x00);
  g_assert_cmpint (cmd[6], ==, 0x00);
  g_assert_cmpint (cmd[7], ==, 0x00);
  g_assert_cmpint (cmd[8], ==, 0x00);
  /* size at bytes 9-12 (LE uint32 = 0x1000) */
  g_assert_cmpint (cmd[9], ==, 0x00);
  g_assert_cmpint (cmd[10], ==, 0x10);
  g_assert_cmpint (cmd[11], ==, 0x00);
  g_assert_cmpint (cmd[12], ==, 0x00);
}

/* ================================================================
 * Regression: Bug #3 — Flash response has 6-byte header
 *
 * After vcsfw_cmd_send strips the 2-byte VCSFW status, the flash
 * response still contains a 6-byte header: [size:4 LE][unknown:2].
 * Actual flash data starts at offset 6. Old code passed the raw
 * response directly to parse_flash(), corrupting the block parsing.
 * This test verifies:
 * 1) The 6-byte header is correctly structured (size field matches)
 * 2) parse_flash works on correctly unwrapped data (offset +6)
 * 3) The raw response differs from the unwrapped payload, proving
 *    that skipping the header is necessary
 * ================================================================ */
static void
test_flash_response_header (void)
{
  /* Build a minimal valid flash image: just an end marker */
  guint8 flash_data[] = { 0xFF, 0xFF, 0x00, 0x00 };

  /* Wrap it in the response header format: [size:4 LE][unk:2][data] */
  guint32 data_size = sizeof (flash_data);
  guint8 response[6 + sizeof (flash_data)];
  FP_WRITE_UINT32_LE (response, data_size);
  response[4] = 0x00;   /* unknown byte 1 */
  response[5] = 0x00;   /* unknown byte 2 */
  memcpy (response + 6, flash_data, sizeof (flash_data));

  /* Verify the unwrap logic: read size from offset 0, skip 6-byte header */
  guint32 read_size = FP_READ_UINT32_LE (response);
  g_assert_cmpuint (read_size, ==, sizeof (flash_data));

  const guint8 *payload = response + 6;
  gsize payload_len = sizeof (response) - 6;
  g_assert_cmpuint (payload_len, ==, sizeof (flash_data));

  /* Parsing the correctly unwrapped data should succeed (end marker only = no keys) */
  ValidityTlsState tls;
  validity_tls_init (&tls);
  GError *error = NULL;
  gboolean result = validity_tls_parse_flash (&tls, payload, payload_len, &error);
  /* Expected: fails because no keys, but NOT because of corrupt block headers */
  g_assert_false (result);
  g_assert_nonnull (error);
  g_assert_true (g_str_has_prefix (error->message, "TLS flash: incomplete key data"));
  g_clear_error (&error);
  validity_tls_free (&tls);

  /* Verify the bug scenario: passing the raw response (with the 6-byte
   * header) gives DIFFERENT data to the parser than the correctly unwrapped
   * payload. The first 4 bytes of the raw response are the LE size field
   * (0x04 0x00 0x00 0x00), which would be misinterpreted as block_id=0x0004
   * (PRIVKEY block with size 0). This is a data corruption — the parser
   * receives wrong input either way, but the key point is that the raw
   * response and the unwrapped payload are NOT the same buffer content. */
  g_assert_cmpuint (sizeof (response), !=, payload_len);
  g_assert_true (memcmp (response, payload, payload_len) != 0);
}

/* ================================================================
 * Regression: Bug #4 — TLS handshake expects raw TLS records
 *
 * parse_server_hello expects raw TLS records starting with a content
 * type byte (0x16 for Handshake). The old code used vcsfw_cmd_send
 * which strips 2 bytes of VCSFW status, corrupting the TLS record.
 * This test verifies that:
 * - A valid TLS Handshake record header is accepted
 * - Data prefixed with a 2-byte VCSFW status is rejected
 * ================================================================ */
static void
test_server_hello_rejects_vcsfw_prefix (void)
{
  /* Build a minimal valid TLS ServerHello record */
  guint8 server_hello_msg[] = {
    /* Handshake message: ServerHello (type 0x02) */
    0x02,               /* type: ServerHello */
    0x00, 0x00, 0x26,   /* length: 38 bytes */
    0x03, 0x03,          /* version 1.2 */
    /* 32 bytes server_random */
    0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08,
    0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f, 0x10,
    0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18,
    0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f, 0x20,
    0x00,                /* session_id length: 0 */
    0xC0, 0x05,          /* cipher suite: 0xC005 */
    0x00,                /* compression: none */
  };

  gsize hs_len = sizeof (server_hello_msg);

  /* Wrap in TLS record: content_type(1) + version(2) + length(2) + body */
  gsize raw_tls_len = 5 + hs_len;
  guint8 *raw_tls = g_malloc (raw_tls_len);
  raw_tls[0] = TLS_CONTENT_HANDSHAKE;   /* 0x16 */
  raw_tls[1] = TLS_VERSION_MAJOR;
  raw_tls[2] = TLS_VERSION_MINOR;
  raw_tls[3] = (hs_len >> 8) & 0xff;
  raw_tls[4] = hs_len & 0xff;
  memcpy (raw_tls + 5, server_hello_msg, hs_len);

  /* Test 1: parse_server_hello with raw TLS — should succeed */
  ValidityTlsState tls;
  validity_tls_init (&tls);
  tls.handshake_hash = g_checksum_new (G_CHECKSUM_SHA256);
  GError *error = NULL;

  gboolean result = validity_tls_parse_server_hello (&tls, raw_tls,
                                                      raw_tls_len, &error);
  g_assert_no_error (error);
  g_assert_true (result);
  /* Verify server_random was properly extracted */
  g_assert_cmpint (tls.server_random[0], ==, 0x01);
  g_assert_cmpint (tls.server_random[31], ==, 0x20);
  validity_tls_free (&tls);

  /* Test 2: Prepend a 2-byte VCSFW status (0x0000) — simulates what
   * vcsfw_cmd_send's cmd_receive_cb would have already STRIPPED.
   * But if the raw recv path is wrong and doesn't strip, the parser
   * gets [0x00, 0x00, 0x16, ...] — first byte 0x00 is not a valid
   * TLS content type, so parsing should behave differently. */
  gsize prefixed_len = 2 + raw_tls_len;
  guint8 *prefixed = g_malloc (prefixed_len);
  prefixed[0] = 0x00;  /* VCSFW status lo */
  prefixed[1] = 0x00;  /* VCSFW status hi */
  memcpy (prefixed + 2, raw_tls, raw_tls_len);

  validity_tls_init (&tls);
  tls.handshake_hash = g_checksum_new (G_CHECKSUM_SHA256);

  result = validity_tls_parse_server_hello (&tls, prefixed, prefixed_len,
                                            &error);
  /* With the 2-byte prefix, the first "record" starts at byte 0:
   * content_type=0x00 is NOT TLS_CONTENT_HANDSHAKE (0x16), so the
   * parser treats it as unknown content and either fails or skips it,
   * and the server_random will NOT match the expected values. */
  if (result)
    {
      /* Even if parsing didn't error, server_random should be wrong */
      gboolean random_ok = (tls.server_random[0] == 0x01 &&
                            tls.server_random[31] == 0x20);
      g_assert_false (random_ok);
    }
  g_clear_error (&error);
  validity_tls_free (&tls);

  g_free (raw_tls);
  g_free (prefixed);
}

/* ================================================================
 * Regression: Bug #5 — Client hello has 0x44 prefix (not VCSFW cmd)
 *
 * TLS handshake messages use 0x44000000 as a 4-byte prefix, NOT a
 * standard VCSFW command byte. This test verifies the prefix and that
 * the TLS record immediately follows (no VCSFW status expected in
 * response).
 * ================================================================ */
static void
test_client_hello_tls_prefix (void)
{
  ValidityTlsState tls;
  validity_tls_init (&tls);

  gsize out_len;
  guint8 *hello = validity_tls_build_client_hello (&tls, &out_len);
  g_assert_nonnull (hello);

  /* Must start with 0x44 0x00 0x00 0x00 (TLS prefix, not VCSFW) */
  g_assert_cmpint (hello[0], ==, 0x44);
  g_assert_cmpint (hello[1], ==, 0x00);
  g_assert_cmpint (hello[2], ==, 0x00);
  g_assert_cmpint (hello[3], ==, 0x00);

  /* Byte 4 must be TLS Handshake content type (0x16) */
  g_assert_cmpint (hello[4], ==, TLS_CONTENT_HANDSHAKE);

  /* Bytes 5-6 must be TLS version 1.2 (0x0303) */
  g_assert_cmpint (hello[5], ==, TLS_VERSION_MAJOR);
  g_assert_cmpint (hello[6], ==, TLS_VERSION_MINOR);

  /* The prefix (0x44) must NOT equal any VCSFW command byte.
   * Specifically, 0x44 != VCSFW_CMD_READ_FLASH (0x40) and
   * is not any known VCSFW command. This proves TLS messages
   * travel on a separate "channel". */
  g_assert_cmpint (hello[0], !=, VCSFW_CMD_GET_VERSION);
  g_assert_cmpint (hello[0], !=, VCSFW_CMD_READ_FLASH);
  g_assert_cmpint (hello[0], !=, VCSFW_CMD_GET_FW_INFO);

  g_free (hello);
  validity_tls_free (&tls);
}

/* ================================================================
 * Main
 * ================================================================ */

int
main (int argc, char *argv[])
{
  g_test_init (&argc, &argv, NULL);

  g_test_add_func ("/validity/tls/prf/deterministic", test_prf_deterministic);
  g_test_add_func ("/validity/tls/prf/output-length", test_prf_output_length);
  g_test_add_func ("/validity/tls/prf/short-output", test_prf_short_output);
  g_test_add_func ("/validity/tls/encrypt/roundtrip",
                   test_encrypt_decrypt_roundtrip);
  g_test_add_func ("/validity/tls/encrypt/block-aligned",
                   test_encrypt_block_aligned);
  g_test_add_func ("/validity/tls/decrypt/invalid", test_decrypt_invalid);
  g_test_add_func ("/validity/tls/psk/derivation", test_psk_derivation);
  g_test_add_func ("/validity/tls/psk/deterministic", test_psk_deterministic);
  g_test_add_func ("/validity/tls/flash/parse-empty", test_flash_parse_empty);
  g_test_add_func ("/validity/tls/flash/parse-truncated",
                   test_flash_parse_truncated);
  g_test_add_func ("/validity/tls/init-free", test_init_free);
  g_test_add_func ("/validity/tls/client-hello", test_build_client_hello);
  g_test_add_func ("/validity/tls/unwrap/invalid", test_unwrap_invalid);

  /* Regression tests for hardware-discovered bugs */
  g_test_add_func ("/validity/tls/regression/flash-parse-needs-psk",
                   test_flash_parse_needs_psk);
  g_test_add_func ("/validity/tls/regression/flash-cmd-format",
                   test_flash_cmd_format);
  g_test_add_func ("/validity/tls/regression/flash-response-header",
                   test_flash_response_header);
  g_test_add_func ("/validity/tls/regression/server-hello-rejects-vcsfw-prefix",
                   test_server_hello_rejects_vcsfw_prefix);
  g_test_add_func ("/validity/tls/regression/client-hello-tls-prefix",
                   test_client_hello_tls_prefix);

  return g_test_run ();
}