philbin 1.0.1

use philbin::careful::Nonce;
use philbin::easy::{
  AssociatedData, AuthTag, AuthTag128, Ciphertext, CiphertextMut, Error, Key,
  Plaintext, PlaintextMut, Result,
};
use std::assert_matches;
use test_utils::AegisTestCase;
use zeroize::Zeroize;

// Takes a provided `philbin::careful` algo submod, test module name,
// auth tag bits and a test case path regex.
// Generates tests that verify the algo against those test cases.
//
// Call example:
//  gen_file_tests!(aegis128l, ietf, 128, "aegis-128l/ietf/*.json");
macro_rules! gen_file_tests_for_tag {
  (
    $algo_mod:ident,     // Should be a `philbin::careful` algo submod
    $test_mod:ident,     // Module name that should group the tests
    $tag_bits:tt,        // 128 or 256, for auth tag num bits
    $files_regex:literal // Should be test case path regex, under "cases"
  ) => {
    pastey::paste! {
      mod [< $test_mod _tag $tag_bits >] {
        use super::*;

        #[rstest::rstest]
        fn [< file_case_tag $tag_bits >] (
          #[base_dir = "../test-utils/src/cases"]
          #[files($files_regex)]
          #[mode = str]
          file_content: &str,
        ) -> anyhow::Result<()> {
          const KEY_BYTES: usize = file_tests::aegis_key_bytes!($algo_mod);
          file_tests::verify_test_json!(
            $algo_mod,
            philbin::easy::[< AuthTag $tag_bits >],
            KEY_BYTES,
            test_utils::from_json(file_content)?
          )
        }
      }
    }
  };
}
pub(crate) use gen_file_tests_for_tag;

// Takes a provided `philbin::careful` algo submod, test module name, and a test
// case path regex.
// Generates tests that verify the algo against those test cases.
// Tests are generated for both 128 bit and 256 bit tag sizes.
//
// Call example:
//  gen_file_tests!(aegis128l, ietf, "aegis-128l/ietf/*.json");
macro_rules! gen_file_tests {
  (
    $algo_mod:ident,     // Should be a `philbin::careful` algo submod
    $test_mod:ident,     // Module name that should group the tests
    $files_regex:literal // Should be test case path regex, under "cases"
  ) => {
    crate::file_tests::gen_file_tests_for_tag!(
      $algo_mod,
      $test_mod,
      128,
      $files_regex
    );
    crate::file_tests::gen_file_tests_for_tag!(
      $algo_mod,
      $test_mod,
      256,
      $files_regex
    );
  };
}
pub(crate) use gen_file_tests;

// Takes an AEGIS submod and returns the number of bytes in the key/nonce for
// that algo.
macro_rules! aegis_key_bytes {
  (aegis128l) => {
    16
  };
  (aegis128x2) => {
    16
  };
  (aegis128x4) => {
    16
  };

  (aegis256) => {
    32
  };
  (aegis256x2) => {
    32
  };
  (aegis256x4) => {
    32
  };
}
pub(crate) use aegis_key_bytes;

// Verifies the provided `algo_mod` AEGIS algo against the provided
// AegisTestCase.
//
// Call example:
//   verify_test_json!(aegis128l, AuthTag128, cases::from_json(file_content)?)
macro_rules! verify_test_json {
  (
    $algo_mod:ident, // Should be a `philbin::careful` algo submod
    $tag_type:ty,    // Should be AuthTag128 or AuthTag256
    $key_bytes:expr,
    $test_case:expr  // Should evaluate to AegisTestCase
  ) => {{
    let enc_in_place_adapted =
      file_tests::encrypt_in_place_adapter::<$tag_type, $key_bytes>(
        philbin::careful::$algo_mod::encrypt_in_place_detached,
      );
    let dec_in_place_adapted =
      file_tests::decrypt_in_place_adapter::<$tag_type, $key_bytes>(
        philbin::careful::$algo_mod::decrypt_in_place_detached,
      );

    let case = $test_case;
    if case.success {
      // *to_slice_detached
      file_tests::verify_encryption::<$tag_type, $key_bytes>(
        philbin::careful::$algo_mod::encrypt_to_slice_detached,
        case.clone(),
      )?;
      file_tests::verify_decryption::<$tag_type, $key_bytes>(
        philbin::careful::$algo_mod::decrypt_to_slice_detached,
        case.clone(),
      )?;

      // *in-place
      file_tests::verify_encryption::<$tag_type, $key_bytes>(
        enc_in_place_adapted,
        case.clone(),
      )?;
      file_tests::verify_decryption::<$tag_type, $key_bytes>(
        dec_in_place_adapted,
        case.clone(),
      )?;
    } else {
      // *to_slice_detached
      file_tests::verify_decryption_failed::<$tag_type, $key_bytes>(
        philbin::careful::$algo_mod::decrypt_to_slice_detached,
        case.clone(),
      )?;

      // *in-place
      file_tests::verify_decryption_failed::<$tag_type, $key_bytes>(
        dec_in_place_adapted,
        case.clone(),
      )?;
    }

    Ok(())
  }};
}
pub(crate) use verify_test_json;

// Adapts `encrypt_in_place` functions to have the API of
// `encrypt_in_place_detached`. Helps with test code reuse.
pub fn encrypt_in_place_adapter<Tag: AuthTag, const KEY_BYTES: usize>(
  func: impl FnOnce(
    PlaintextMut,
    AssociatedData,
    &Key<KEY_BYTES>,
    Nonce<KEY_BYTES>,
  ) -> Result<Tag>,
) -> impl FnOnce(
  Plaintext,
  AssociatedData,
  &Key<KEY_BYTES>,
  Nonce<KEY_BYTES>,
  CiphertextMut,
) -> Result<Tag> {
  |plaintext, associated_data, key, nonce, mut ciphertext| {
    // Since we always need to copy the `plaintext` bytes into the `buffer` for
    // in-place processing, we need a replacement `buffer` since `ciphertext` is
    // zero-length when AegisTestCase::success is false.
    let mut buffer = vec![0; plaintext.as_bytes().len()];
    buffer.copy_from_slice(plaintext.as_bytes());
    let pt_buffer = PlaintextMut::new(&mut buffer);
    func(pt_buffer, associated_data, key, nonce)
      .inspect(|_| ciphertext.as_bytes_mut().copy_from_slice(&buffer))
  }
}

// Adapts `decrypt_in_place` functions to have the API of
// `decrypt_in_place_detached`. Helps with test code reuse.
pub fn decrypt_in_place_adapter<Tag: AuthTag, const KEY_BYTES: usize>(
  func: impl FnOnce(
    CiphertextMut,
    &Tag,
    AssociatedData,
    &Key<KEY_BYTES>,
    Nonce<KEY_BYTES>,
  ) -> Result<()>,
) -> impl FnOnce(
  Ciphertext,
  &Tag,
  AssociatedData,
  &Key<KEY_BYTES>,
  Nonce<KEY_BYTES>,
  PlaintextMut,
) -> Result<()> {
  |ciphertext, tag, associated_data, key, nonce, mut plaintext| {
    // Since we always need to copy the `ciphertext` bytes into the `buffer` for
    // in-place processing, we need a replacement `buffer` since `plaintext` is
    // zero-length when AegisTestCase::success is false.
    let mut buffer = vec![0; ciphertext.as_bytes().len()];
    buffer.copy_from_slice(ciphertext.as_bytes());
    let ct_buffer = CiphertextMut::new(&mut buffer);
    let result = func(ct_buffer, tag, associated_data, key, nonce);
    if result.is_ok() {
      plaintext.as_bytes_mut().copy_from_slice(&buffer);
    } else {
      // Still must zeroize the right parts of plaintext on decryption failure
      plaintext
        .as_bytes_mut()
        .get_mut(..ciphertext.as_bytes().len())
        .unwrap()
        .zeroize();
    }
    result
  }
}

// NOTE: Use `verify_test_json` instead of calling this directly.
pub fn verify_encryption<Tag: AuthTag, const KEY_BYTES: usize>(
  func: impl FnOnce(
    Plaintext,
    AssociatedData,
    &Key<KEY_BYTES>,
    Nonce<KEY_BYTES>,
    CiphertextMut,
  ) -> Result<Tag>,
  expected: AegisTestCase,
) -> anyhow::Result<()> {
  assert!(expected.success);

  let mut actual_ciphertext = vec![0; expected.plaintext.len()];
  let actual_tag: Tag = func(
    Plaintext::new(&expected.plaintext),
    AssociatedData::new(&expected.associated_data),
    &Key::<KEY_BYTES>::from_bytes(expected.key)?,
    Nonce::<KEY_BYTES>::new(*expected.nonce.as_array().unwrap())?,
    CiphertextMut::new(&mut actual_ciphertext),
  )?;
  assert_eq!(
    expected.ciphertext,
    actual_ciphertext[..expected.ciphertext.len()],
    "actual ciphertext did not match; case notes: {}",
    expected.notes,
  );

  if size_of::<Tag>() == size_of::<AuthTag128>() {
    assert_eq!(
      &expected.tag128,
      actual_tag.as_ref(),
      "actual tag did not match; case notes: {}",
      expected.notes,
    );
  } else {
    assert_eq!(
      &expected.tag256,
      actual_tag.as_ref(),
      "actual tag did not match; case notes: {}",
      expected.notes,
    );
  }

  Ok(())
}

// NOTE: Use `verify_test_json` instead of calling this directly.
pub fn verify_decryption<Tag: AuthTag, const KEY_BYTES: usize>(
  func: impl FnOnce(
    Ciphertext,
    &Tag,
    AssociatedData,
    &Key<KEY_BYTES>,
    Nonce<KEY_BYTES>,
    PlaintextMut,
  ) -> Result<()>,
  expected: AegisTestCase,
) -> anyhow::Result<()> {
  assert!(expected.success);

  let mut actual_plaintext = vec![0; expected.ciphertext.len()];
  let mut expected_tag = Tag::default();
  if size_of::<Tag>() == size_of::<AuthTag128>() {
    expected_tag.as_mut().copy_from_slice(&expected.tag128);
  } else {
    expected_tag.as_mut().copy_from_slice(&expected.tag256);
  }

  func(
    Ciphertext::new(&expected.ciphertext),
    &expected_tag,
    AssociatedData::new(&expected.associated_data),
    &Key::<KEY_BYTES>::from_bytes(expected.key)?,
    Nonce::<KEY_BYTES>::new(*expected.nonce.as_array().unwrap())?,
    PlaintextMut::new(&mut actual_plaintext),
  )?;
  assert_eq!(
    expected.plaintext, actual_plaintext,
    "decrypted plaintext did not match; case notes: {}",
    expected.notes,
  );

  Ok(())
}

// NOTE: Use `verify_test_json` instead of calling this directly.
pub fn verify_decryption_failed<Tag: AuthTag, const KEY_BYTES: usize>(
  func: impl FnOnce(
    Ciphertext,
    &Tag,
    AssociatedData,
    &Key<KEY_BYTES>,
    Nonce<KEY_BYTES>,
    PlaintextMut,
  ) -> Result<()>,
  expected: AegisTestCase,
) -> anyhow::Result<()> {
  assert!(!expected.success);

  // Intentionally longer than necessary (and set to ones, not zeros) so we
  // can verify how much iz zeroed out on tag auth failure.
  let mut actual_plaintext = vec![1; expected.ciphertext.len() + 100];

  let mut expected_tag = Tag::default();
  if size_of::<Tag>() == size_of::<AuthTag128>() {
    expected_tag.as_mut().copy_from_slice(&expected.tag128);
  } else {
    expected_tag.as_mut().copy_from_slice(&expected.tag256);
  }

  assert_matches!(
    func(
      Ciphertext::new(&expected.ciphertext),
      &expected_tag,
      AssociatedData::new(&expected.associated_data),
      &Key::<KEY_BYTES>::from_bytes(expected.key)?,
      Nonce::<KEY_BYTES>::new(*expected.nonce.as_array().unwrap())?,
      PlaintextMut::new(&mut actual_plaintext),
    ),
    Err(Error::AuthTagInvalid),
    "decryption did not fail; case notes: {}",
    expected.notes
  );

  // The first `ciphertext.len()` amount of bytes in `plaintext` should be
  // zeroed out.
  assert_eq!(
    &vec![0; expected.ciphertext.len()],
    &actual_plaintext[..expected.ciphertext.len()],
    "plaintext bytes were not zeroed out; case notes: {}",
    expected.notes
  );
  // The remainder are still set to ones.
  assert_eq!(
    &vec![1; actual_plaintext.len() - expected.ciphertext.len()],
    &actual_plaintext[expected.ciphertext.len()..],
    "unrelated plaintext bytes were modified; case notes: {}",
    expected.notes
  );

  Ok(())
}