// Copyright 2022 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

extern crate alloc;

use alloc::vec;
use crypto_provider::aes::BLOCK_SIZE;
use crypto_provider_rustcrypto::RustCrypto;
use ldt::{DefaultPadder, Ldt, LdtError, LdtKey, Padder, Swap, XorPadder};
use xts_aes::{XtsAes128, XtsAes128Key};

#[test]
fn normal_pad_empty() {
    let padder = DefaultPadder::default();
    let tweak: xts_aes::Tweak =
        <DefaultPadder as Padder<16, xts_aes::XtsAes128<RustCrypto>>>::pad_tweak(&padder, &[]);

    let bytes = tweak.le_bytes();

    // leading 1 bit
    let mut expected = [0; 16];
    expected[0] = 0x80;
    assert_eq!(expected, bytes);
}

#[test]
fn normal_pad_one_byte() {
    let padder = DefaultPadder::default();
    let tweak: xts_aes::Tweak =
        <DefaultPadder as Padder<16, xts_aes::XtsAes128<RustCrypto>>>::pad_tweak(&padder, &[0x81]);

    let bytes = tweak.le_bytes();

    // 1 bit after original byte
    let mut expected = [0; 16];
    expected[0] = 0x81;
    expected[1] = 0x80;
    assert_eq!(expected, bytes);
}

#[test]
fn normal_pad_max_len() {
    let padder = DefaultPadder::default();
    let input = [0x99; 15];
    let tweak: xts_aes::Tweak =
        <DefaultPadder as Padder<16, XtsAes128<RustCrypto>>>::pad_tweak(&padder, &input);

    let bytes = tweak.le_bytes();

    // 1 bit after original bytes
    let mut expected = [0x99; 16];
    expected[15] = 0x80;
    assert_eq!(expected, bytes);
}

#[test]
#[should_panic]
fn normal_pad_too_big_panics() {
    let padder = DefaultPadder::default();
    let input = [0x99; 16];
    <DefaultPadder as Padder<16, xts_aes::XtsAes128<RustCrypto>>>::pad_tweak(&padder, &input);
}

#[test]
fn xor_pad_empty() {
    let padder = [0x24; BLOCK_SIZE].into();
    let tweak: xts_aes::Tweak = <XorPadder<BLOCK_SIZE> as Padder<
        BLOCK_SIZE,
        xts_aes::XtsAes128<RustCrypto>,
    >>::pad_tweak(&padder, &[]);

    let bytes = tweak.le_bytes();

    // leading 1 bit
    let mut expected = [0x24; BLOCK_SIZE];
    expected[0] = 0x24 ^ 0x80;
    assert_eq!(expected, bytes);
}

#[test]
fn xor_pad_one_byte() {
    let padder = [0x24; BLOCK_SIZE].into();
    let tweak: xts_aes::Tweak = <XorPadder<BLOCK_SIZE> as Padder<
        BLOCK_SIZE,
        XtsAes128<RustCrypto>,
    >>::pad_tweak(&padder, &[0x81]);

    let bytes = tweak.le_bytes();

    // 1 bit after original byte
    let mut expected = [0x24; BLOCK_SIZE];
    expected[0] = 0x81 ^ 0x24;
    expected[1] = 0x80 ^ 0x24;
    assert_eq!(expected, bytes);
}

#[test]
fn xor_pad_max_len() {
    let padder = [0x24; BLOCK_SIZE].into();
    let input = [0x99; 15];
    let tweak: xts_aes::Tweak = <XorPadder<BLOCK_SIZE> as Padder<
        BLOCK_SIZE,
        XtsAes128<RustCrypto>,
    >>::pad_tweak(&padder, &input);

    let bytes = tweak.le_bytes();

    // 1 bit after original bytes
    let mut expected = [0x99 ^ 0x24; BLOCK_SIZE];
    expected[15] = 0x80 ^ 0x24;
    assert_eq!(expected, bytes);
}

#[test]
#[should_panic]
fn xor_pad_too_big_panics() {
    let padder = [0x24; BLOCK_SIZE].into();
    // need 1 byte for padding, and 2 more for salt
    let input = [0x99; 16];
    <XorPadder<BLOCK_SIZE> as Padder<BLOCK_SIZE, XtsAes128<RustCrypto>>>::pad_tweak(
        &padder, &input,
    );
}

#[test]
fn encrypt_too_short_err() {
    do_length_check(7, |ldt, payload| {
        ldt.encrypt(payload, &DefaultPadder::default())
    })
}

#[test]
fn encrypt_too_long_err() {
    do_length_check(40, |ldt, payload| {
        ldt.encrypt(payload, &DefaultPadder::default())
    })
}
#[test]
fn decrypt_too_short_err() {
    do_length_check(7, |ldt, payload| {
        ldt.decrypt(payload, &DefaultPadder::default())
    })
}
#[test]
fn decrypt_too_long_err() {
    do_length_check(40, |ldt, payload| {
        ldt.decrypt(payload, &DefaultPadder::default())
    })
}

fn do_length_check<
    F: Fn(Ldt<{ BLOCK_SIZE }, XtsAes128<RustCrypto>, Swap>, &mut [u8]) -> Result<(), LdtError>,
>(
    payload_len: usize,
    ldt_op: F,
) {
    let ldt = Ldt::new(&LdtKey::<XtsAes128Key>::from_concatenated(&[0u8; 64]));

    let mut payload = vec![0; payload_len];
    assert_eq!(
        Err(LdtError::InvalidLength(payload_len)),
        ldt_op(ldt, &mut payload)
    );
}