// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use super::*;
use crate::varint::{VarInt, MAX_VARINT_VALUE};
use bolero::{check, generator::*};

static BYTES: &[u8] = &[42u8; 9000];

#[derive(Copy, Clone, Debug, TypeGenerator)]
enum Op {
    Write {
        offset: VarInt,
        #[generator(0..BYTES.len())]
        len: usize,
    },
    Pop {
        watermark: Option<u16>,
    },
}

#[test]
#[cfg_attr(miri, ignore)] // This test is too expensive for miri to complete in a reasonable amount of time
fn model_test() {
    check!().with_type::<Vec<Op>>().for_each(|ops| {
        let mut buffer = ReceiveBuffer::new();
        for op in ops {
            match *op {
                Op::Write { offset, len } => {
                    let _ = buffer.write_at(offset, &BYTES[..len]);
                }
                Op::Pop { watermark } => {
                    if let Some(watermark) = watermark {
                        if let Some(chunk) = buffer.pop_watermarked(watermark as _) {
                            assert!(chunk.len() <= watermark as usize);
                        }
                    } else if let Some(chunk) = buffer.pop() {
                        assert!(!chunk.is_empty(), "popped chunks should never be empty");
                    }
                }
            }
        }
    })
}

#[test]
#[cfg_attr(miri, ignore)] // This test is too expensive for miri to complete in a reasonable amount of time
fn write_and_pop() {
    let mut buffer = ReceiveBuffer::new();
    let mut offset = VarInt::default();
    let mut popped_bytes = 0;
    for _ in 0..10000 {
        buffer.write_at(offset, BYTES).unwrap();
        offset += BYTES.len();
        while let Some(chunk) = buffer.pop() {
            popped_bytes += chunk.len();
        }
    }
    assert_eq!(offset, popped_bytes);
}

fn new_receive_buffer() -> ReceiveBuffer {
    let buffer = ReceiveBuffer::new();
    assert_eq!(buffer.len(), 0);
    buffer
}

#[test]
fn gap_replacement_test() {
    let mut buffer = new_receive_buffer();

    assert!(buffer.write_at(4u32.into(), &[4, 5, 6, 7]).is_ok());
    assert_eq!(0, buffer.len());

    assert!(buffer.write_at(0u32.into(), &[0, 1, 2, 3]).is_ok());
    assert_eq!(8, buffer.len());

    assert_eq!(&[0u8, 1, 2, 3, 4, 5, 6, 7], &buffer.pop().unwrap()[..]);
    assert_eq!(None, buffer.pop());
}

#[test]
fn gap_subset_test() {
    let mut buffer = new_receive_buffer();

    assert!(buffer.write_at(4u32.into(), &[4, 5, 6, 7]).is_ok());
    assert_eq!(0, buffer.len());

    assert!(buffer.write_at(2u32.into(), &[2]).is_ok());
    assert_eq!(0, buffer.len());
    assert!(buffer.write_at(0u32.into(), &[0, 1, 2, 3]).is_ok());
    assert_eq!(8, buffer.len());

    assert_eq!(&[0u8, 1, 2, 3, 4, 5, 6, 7], &buffer.pop().unwrap()[..]);
    assert!(buffer.is_empty());
}

#[test]
fn gap_beginning_test() {
    let mut buffer = new_receive_buffer();

    assert!(buffer.write_at(0u32.into(), &[0, 1, 2, 3]).is_ok());
    assert_eq!(4, buffer.len());

    assert_eq!(&[0u8, 1, 2, 3], &buffer.pop().unwrap()[..]);
    assert!(buffer.is_empty());
}

#[test]
fn gap_end_test() {
    let mut buffer = new_receive_buffer();

    assert!(buffer.write_at(4u32.into(), &[4, 5, 6, 7]).is_ok());
    assert_eq!(0, buffer.len());

    assert!(buffer.write_at(3u32.into(), &[3]).is_ok());
    assert_eq!(0, buffer.len());

    assert!(buffer.write_at(0u32.into(), &[0, 1, 2]).is_ok());
    assert_eq!(8, buffer.len());

    assert_eq!(&[0u8, 1, 2, 3, 4, 5, 6, 7], &buffer.pop().unwrap()[..]);
    assert!(buffer.is_empty());
}

#[test]
fn gap_overlap_test() {
    let mut buffer = new_receive_buffer();

    assert!(buffer.write_at(4u32.into(), &[4, 5, 6, 7]).is_ok());
    assert_eq!(0, buffer.len());

    assert!(buffer.write_at(2u32.into(), &[2]).is_ok());
    assert_eq!(0, buffer.len());

    assert!(buffer.write_at(1u32.into(), &[1, 2, 3]).is_ok());
    assert_eq!(0, buffer.len());

    assert!(buffer.write_at(0u32.into(), &[0]).is_ok());
    assert_eq!(8, buffer.len());

    assert_eq!(&[0u8, 1, 2, 3, 4, 5, 6, 7], &buffer.pop().unwrap()[..]);
    assert!(buffer.is_empty());
}

#[test]
fn gap_later_test() {
    let mut buffer = new_receive_buffer();

    assert!(buffer.write_at(4u32.into(), &[4, 5, 6, 7]).is_ok());
    assert_eq!(0, buffer.len());

    assert!(buffer.write_at(2u32.into(), &[2]).is_ok());
    assert_eq!(0, buffer.len());

    assert!(buffer.write_at(3u32.into(), &[3]).is_ok());
    assert_eq!(0, buffer.len());

    assert!(buffer.write_at(0u32.into(), &[0, 1]).is_ok());
    assert_eq!(8, buffer.len());

    assert_eq!(&[0u8, 1, 2, 3, 4, 5, 6, 7], &buffer.pop().unwrap()[..]);
    assert!(buffer.is_empty());
}

#[test]
fn gap_multiple_subset_test() {
    let mut buffer = new_receive_buffer();

    assert!(buffer.write_at(4u32.into(), &[4, 5, 6, 7]).is_ok());
    assert_eq!(0, buffer.len());

    assert!(buffer.write_at(1u32.into(), &[1]).is_ok());
    assert_eq!(0, buffer.len());

    assert!(buffer.write_at(3u32.into(), &[3]).is_ok());
    assert_eq!(0, buffer.len());

    assert!(buffer.write_at(0u32.into(), &[0, 1, 2, 3]).is_ok());
    assert_eq!(8, buffer.len());

    assert_eq!(&[0u8, 1, 2, 3, 4, 5, 6, 7], &buffer.pop().unwrap()[..]);
    assert!(buffer.is_empty());
}

#[test]
fn gap_partial_overlay_test() {
    let mut buffer = new_receive_buffer();

    assert!(buffer.write_at(4u32.into(), &[4, 5, 6, 7]).is_ok());
    assert_eq!(0, buffer.len());

    assert!(buffer.write_at(1u32.into(), &[1]).is_ok());
    assert_eq!(0, buffer.len());

    assert!(buffer.write_at(0u32.into(), &[0, 1, 2]).is_ok());
    assert_eq!(3, buffer.len());

    assert_eq!(&[0u8, 1, 2], &buffer.pop().unwrap()[..]);
    assert_eq!(0, buffer.len());

    assert!(buffer.write_at(3u32.into(), &[3]).is_ok());
    assert_eq!(5, buffer.len());

    assert_eq!(&[3u8, 4, 5, 6, 7], &buffer.pop().unwrap()[..]);
    assert!(buffer.is_empty());
}

#[test]
fn gap_partial_stale_test() {
    let mut buffer = new_receive_buffer();

    assert!(buffer.write_at(4u32.into(), &[4, 5]).is_ok());
    assert_eq!(0, buffer.len());

    assert!(buffer.write_at(0u32.into(), &[0, 1]).is_ok());
    assert_eq!(2, buffer.len());

    assert_eq!(&[0u8, 1], &buffer.pop().unwrap()[..]);
    assert_eq!(0, buffer.len());

    assert_eq!(None, buffer.pop());
    assert_eq!(0, buffer.len());

    assert!(buffer.write_at(0u32.into(), &[0, 1, 2, 3]).is_ok());
    assert_eq!(4, buffer.len());

    assert_eq!(&[2, 3, 4, 5], &buffer.pop().unwrap()[..]);
    assert!(buffer.is_empty());
}

#[test]
fn chunk_duplicate_test() {
    let mut buffer = new_receive_buffer();

    assert!(buffer.write_at(0u32.into(), &[0, 1, 2, 3]).is_ok());
    assert_eq!(4, buffer.len());

    assert!(buffer.write_at(0u32.into(), &[10, 11, 12, 13]).is_ok()); // exact match
    assert!(buffer.write_at(0u32.into(), &[20, 21]).is_ok()); // beginning
    assert!(buffer.write_at(1u32.into(), &[31, 32]).is_ok()); // middle
    assert!(buffer.write_at(2u32.into(), &[42, 43]).is_ok()); // end

    assert_eq!(4, buffer.len());
    assert_eq!(&[0, 1, 2, 3], &buffer.pop().unwrap()[..]);
}

#[test]
fn chunk_superset_test() {
    let mut buffer = new_receive_buffer();

    assert!(buffer.write_at(0u32.into(), &[0]).is_ok());
    assert_eq!(1, buffer.len());

    assert!(buffer.write_at(1u32.into(), &[1]).is_ok());
    assert_eq!(2, buffer.len());

    assert!(buffer.write_at(2u32.into(), &[2]).is_ok());
    assert_eq!(3, buffer.len());

    assert!(buffer.write_at(3u32.into(), &[3]).is_ok());
    assert_eq!(4, buffer.len());

    assert!(buffer.write_at(0u32.into(), &[0, 1, 2, 3, 4]).is_ok());
    assert_eq!(5, buffer.len());

    assert_eq!(&[0u8, 1, 2, 3, 4], &buffer.pop().unwrap()[..]);
    assert!(buffer.is_empty());
}

#[test]
fn chunk_stale_test() {
    let mut buffer = new_receive_buffer();

    assert!(buffer.write_at(0u32.into(), &[0, 1, 2, 3]).is_ok());
    assert_eq!(4, buffer.len());

    assert_eq!(&[0u8, 1, 2, 3], &buffer.pop().unwrap()[..]);
    assert!(buffer.is_empty());

    assert!(buffer.write_at(4u32.into(), &[4, 5, 6, 7]).is_ok());
    assert_eq!(4, buffer.len());

    // ignore stale data
    assert!(buffer.write_at(0u32.into(), &[0, 1, 2, 3]).is_ok());
    assert_eq!(4, buffer.len());

    assert_eq!(&[4, 5, 6, 7], &buffer.pop().unwrap()[..]);
    assert!(buffer.is_empty());
}

#[test]
fn chunk_partial_smaller_before_test() {
    let mut buffer = new_receive_buffer();

    assert!(buffer.write_at(4u32.into(), &[4, 5, 6, 7]).is_ok());
    assert_eq!(0, buffer.len());

    assert!(buffer.write_at(2u32.into(), &[2, 3, 4, 5]).is_ok());
    assert_eq!(0, buffer.len());

    assert!(buffer.write_at(0u32.into(), &[0, 1]).is_ok());
    assert_eq!(8, buffer.len());

    assert_eq!(&[0u8, 1, 2, 3, 4, 5, 6, 7], &buffer.pop().unwrap()[..]);
    assert!(buffer.is_empty());
}

#[test]
fn chunk_partial_larger_before_test() {
    let mut buffer = new_receive_buffer();

    assert!(buffer.write_at(4u32.into(), &[4, 5, 6]).is_ok());
    assert_eq!(0, buffer.len());

    assert!(buffer.write_at(1u32.into(), &[1, 2, 3, 4, 5]).is_ok());
    assert_eq!(0, buffer.len());

    assert!(buffer.write_at(0u32.into(), &[0, 2]).is_ok());
    assert_eq!(7, buffer.len());

    assert_eq!(&[0u8, 1, 2, 3, 4, 5, 6], &buffer.pop().unwrap()[..]);
}

#[test]
fn chunk_partial_smaller_after_test() {
    let mut buffer = new_receive_buffer();

    assert!(buffer.write_at(0u32.into(), &[0, 1, 2, 3]).is_ok());
    assert_eq!(4, buffer.len());

    assert!(buffer.write_at(3u32.into(), &[3, 4]).is_ok());
    assert_eq!(5, buffer.len());

    assert_eq!(&[0u8, 1, 2, 3, 4], &buffer.pop().unwrap()[..]);
    assert!(buffer.is_empty());
}

#[test]
fn chunk_partial_larger_after_test() {
    let mut buffer = new_receive_buffer();

    assert!(buffer.write_at(0u32.into(), &[0, 1]).is_ok());
    assert_eq!(2, buffer.len());

    assert!(buffer.write_at(1u32.into(), &[1, 2, 3]).is_ok());
    assert_eq!(4, buffer.len());

    assert_eq!(&[0u8, 1, 2, 3], &buffer.pop().unwrap()[..]);
    assert!(buffer.is_empty());
}

#[test]
#[allow(clippy::cognitive_complexity)] // several operations are needed to get the buffer in the desired state
fn write_and_read_buffer() {
    let mut buf = ReceiveBuffer::new();

    assert_eq!(0, buf.len());
    assert!(buf.is_empty());
    assert_eq!(None, buf.pop());

    buf.write_at(0u32.into(), &[0, 1, 2, 3]).unwrap();
    assert_eq!(4, buf.len());
    assert!(!buf.is_empty());
    assert_eq!(&[0u8, 1, 2, 3], &*buf.pop().unwrap());
    assert_eq!(0, buf.len());
    assert!(buf.is_empty());
    assert_eq!(None, buf.pop());

    buf.write_at(0u32.into(), &[0, 1, 2, 3]).unwrap();
    assert_eq!(0, buf.len());
    assert_eq!(None, buf.pop());

    buf.write_at(2u32.into(), &[2, 3, 4, 5, 6]).unwrap();
    assert_eq!(3, buf.len());
    assert_eq!(&[4, 5, 6], &*buf.pop().unwrap());

    // Create a gap
    buf.write_at(10u32.into(), &[10, 11, 12]).unwrap();
    assert_eq!(0, buf.len());
    assert_eq!(None, buf.pop());

    // Fill the gap
    buf.write_at(7u32.into(), &[7, 8, 9]).unwrap();
    assert_eq!(6, buf.len());
    assert_eq!(&[7, 8, 9, 10, 11, 12], &*buf.pop().unwrap());
    assert_eq!(0, buf.len());
    assert_eq!(None, buf.pop());

    // Create another gap
    buf.write_at(13u32.into(), &[13]).unwrap();
    assert_eq!(1, buf.len());

    buf.write_at(16u32.into(), &[16, 17, 18]).unwrap();
    assert_eq!(1, buf.len());

    // Overfill the gap
    buf.write_at(12u32.into(), &[12, 13, 14, 15, 16, 17, 18, 19, 20])
        .unwrap();
    assert_eq!(8, buf.len());
    assert_eq!(&[13, 14, 15, 16, 17, 18, 19, 20], &*buf.pop().unwrap());
    assert_eq!(0, buf.len());
    assert_eq!(None, buf.pop());
}

#[test]
fn fill_preallocated_gaps() {
    let mut buf: ReceiveBuffer = ReceiveBuffer::new();

    buf.write_at((MIN_BUFFER_ALLOCATION_SIZE as u32 + 2).into(), &[42, 45])
        .unwrap();
    assert_eq!(0, buf.len());
    assert_eq!(None, buf.pop());

    let mut first_chunk = [0u8; MIN_BUFFER_ALLOCATION_SIZE - 2];
    for (idx, b) in first_chunk.iter_mut().enumerate() {
        *b = idx as u8;
    }

    buf.write_at(0u32.into(), &first_chunk).unwrap();
    assert_eq!(MIN_BUFFER_ALLOCATION_SIZE - 2, buf.len());

    // Overfill gap
    buf.write_at(
        (MIN_BUFFER_ALLOCATION_SIZE as u32 - 4).into(),
        &[91, 92, 93, 94, 95, 96, 97, 98, 99],
    )
    .unwrap();
    assert_eq!(MIN_BUFFER_ALLOCATION_SIZE + 5, buf.len());

    let chunk = buf.pop().unwrap();
    assert_eq!(MIN_BUFFER_ALLOCATION_SIZE, chunk.len());
    let mut expected = 0;
    for b in &chunk[..MIN_BUFFER_ALLOCATION_SIZE - 2] {
        assert_eq!(expected, *b);
        expected = expected.wrapping_add(1);
    }
    assert_eq!(93, chunk[MIN_BUFFER_ALLOCATION_SIZE - 2]);
    assert_eq!(94, chunk[MIN_BUFFER_ALLOCATION_SIZE - 1]);

    let chunk = buf.pop().unwrap();
    assert_eq!(5, chunk.len());
    assert_eq!(95, chunk[0]);
    assert_eq!(96, chunk[1]);
    assert_eq!(42, chunk[2]);
    assert_eq!(45, chunk[3]);
    assert_eq!(99, chunk[4]);

    assert_eq!(None, buf.pop());
}

#[test]
fn create_and_fill_large_gaps() {
    let mut buf = ReceiveBuffer::new();
    // This creates 3 full buffer gaps of full allocation ranges
    buf.write_at(
        (MIN_BUFFER_ALLOCATION_SIZE as u32 * 3 + 2).into(),
        &[7, 8, 9],
    )
    .unwrap();
    assert_eq!(0, buf.len());

    // Insert something in the middle, which still leaves us with 2 gaps
    buf.write_at(
        (MIN_BUFFER_ALLOCATION_SIZE as u32 + 1).into(),
        &[10, 11, 12],
    )
    .unwrap();
    assert_eq!(0, buf.len());

    // Close the last gap
    buf.write_at(
        (MIN_BUFFER_ALLOCATION_SIZE as u32 * 2 + 1).into(),
        &[14, 15, 16],
    )
    .unwrap();
    assert_eq!(0, buf.len());

    // Close the first gap
    buf.write_at(4u32.into(), &[4, 5, 6, 7]).unwrap();
    assert_eq!(0, buf.len());

    // Now fill all the preallocated chunks with data
    let data = [0u8; MIN_BUFFER_ALLOCATION_SIZE * 3 + 20];
    buf.write_at(0u32.into(), &data).unwrap();
    assert_eq!(MIN_BUFFER_ALLOCATION_SIZE * 3 + 20, buf.len());

    fn assert_zero_content(slice: &[u8]) {
        for b in slice {
            assert_eq!(0, *b);
        }
    }

    // Check the first allocation range
    let chunk = buf.pop().unwrap();
    assert_eq!(MIN_BUFFER_ALLOCATION_SIZE, chunk.len());
    assert_zero_content(&chunk[..4]);
    assert_eq!(&[4, 5, 6, 7], &chunk[4..8]);
    assert_zero_content(&chunk[8..MIN_BUFFER_ALLOCATION_SIZE]);

    // Check the second allocation range
    let chunk = buf.pop().unwrap();
    assert_eq!(MIN_BUFFER_ALLOCATION_SIZE, chunk.len());
    assert_zero_content(&chunk[..1]);
    assert_eq!(&[10, 11, 12], &chunk[1..4]);
    assert_zero_content(&chunk[4..MIN_BUFFER_ALLOCATION_SIZE]);

    // Check the third allocation range
    let chunk = buf.pop().unwrap();
    assert_eq!(MIN_BUFFER_ALLOCATION_SIZE, chunk.len());
    assert_zero_content(&chunk[..1]);
    assert_eq!(&[14, 15, 16], &chunk[1..4]);
    assert_zero_content(&chunk[4..MIN_BUFFER_ALLOCATION_SIZE]);

    // Check the forth allocation range
    let chunk = buf.pop().unwrap();
    assert_eq!(20, chunk.len());
    assert_zero_content(&chunk[..2]);
    assert_eq!(&[7, 8, 9], &chunk[2..5]);
    assert_zero_content(&chunk[5..20]);

    // No more data should be available
    assert_eq!(0, buf.len());
    assert_eq!(None, buf.pop());
}

#[test]
fn ignore_already_consumed_data() {
    let mut buf = ReceiveBuffer::new();
    buf.write_at(0u32.into(), &[0, 1, 2, 3]).unwrap();
    assert_eq!(0, buf.consumed_len());
    assert_eq!(&[0u8, 1, 2, 3], &*buf.pop().unwrap());
    assert_eq!(4, buf.consumed_len());

    // Add partially consumed data
    buf.write_at(2u32.into(), &[2, 3, 4]).unwrap();
    assert_eq!(4, buf.consumed_len());
    assert_eq!(&[4], &*buf.pop().unwrap());
    assert_eq!(5, buf.consumed_len());

    // Add data which had been fully consumed before
    for start_offset in 0..3 {
        buf.write_at((start_offset as u32).into(), &[0, 1, 2, 3][start_offset..])
            .unwrap();
        assert_eq!(0, buf.len());
        assert_eq!(None, buf.pop());
        assert_eq!(5, buf.consumed_len());
    }
}

#[test]
fn merge_right() {
    let mut buf = ReceiveBuffer::new();
    buf.write_at(4u32.into(), &[4, 5, 6]).unwrap();
    buf.write_at(0u32.into(), &[0, 1, 2, 3]).unwrap();
    assert_eq!(7, buf.len());
    assert_eq!(&[0u8, 1, 2, 3, 4, 5, 6], &*buf.pop().unwrap());
}

#[test]
fn merge_left() {
    let mut buf = ReceiveBuffer::new();
    buf.write_at(0u32.into(), &[0, 1, 2, 3]).unwrap();
    buf.write_at(4u32.into(), &[4, 5, 6]).unwrap();
    assert_eq!(7, buf.len());
    assert_eq!(&[0u8, 1, 2, 3, 4, 5, 6], &*buf.pop().unwrap());
}

#[test]
fn merge_both_sides() {
    let mut buf = ReceiveBuffer::new();
    // Create gaps on all sides, and merge them later
    buf.write_at(4u32.into(), &[4, 5]).unwrap();
    buf.write_at(8u32.into(), &[8, 9]).unwrap();
    buf.write_at(6u32.into(), &[6, 7]).unwrap();
    buf.write_at(0u32.into(), &[0, 1, 2, 3]).unwrap();
    assert_eq!(10, buf.len());
    assert_eq!(&[0u8, 1, 2, 3, 4, 5, 6, 7, 8, 9], &*buf.pop().unwrap());
}

#[test]
fn do_not_merge_across_allocations_right() {
    let mut buf = ReceiveBuffer::new();
    let mut data_left = [0u8; MIN_BUFFER_ALLOCATION_SIZE];
    let mut data_right = [0u8; MIN_BUFFER_ALLOCATION_SIZE];

    data_left[MIN_BUFFER_ALLOCATION_SIZE - 1] = 13;
    data_right[0] = 72;

    buf.write_at((MIN_BUFFER_ALLOCATION_SIZE as u32).into(), &data_right[..])
        .unwrap();
    buf.write_at(0u32.into(), &data_left[..]).unwrap();
    assert_eq!(2 * MIN_BUFFER_ALLOCATION_SIZE, buf.len());

    assert_eq!(&data_left[..], &*buf.pop().unwrap());
    assert_eq!(&data_right[..], &*buf.pop().unwrap());
}

#[test]
fn do_not_merge_across_allocations_left() {
    let mut buf = ReceiveBuffer::new();
    let mut data_left = [0u8; MIN_BUFFER_ALLOCATION_SIZE];
    let mut data_right = [0u8; MIN_BUFFER_ALLOCATION_SIZE];

    data_left[MIN_BUFFER_ALLOCATION_SIZE - 1] = 13;
    data_right[0] = 72;

    buf.write_at(0u32.into(), &data_left[..]).unwrap();
    buf.write_at((MIN_BUFFER_ALLOCATION_SIZE as u32).into(), &data_right[..])
        .unwrap();
    assert_eq!(2 * MIN_BUFFER_ALLOCATION_SIZE, buf.len());

    assert_eq!(&data_left[..], &*buf.pop().unwrap());
    assert_eq!(&data_right[..], &*buf.pop().unwrap());
}

#[test]
fn reset_buffer() {
    let mut buf = ReceiveBuffer::new();
    buf.write_at(2u32.into(), &[2, 3]).unwrap();
    buf.write_at(0u32.into(), &[0, 1]).unwrap();
    assert_eq!(4, buf.len());
    buf.pop().unwrap();
    assert_eq!(0, buf.len());

    buf.write_at(4u32.into(), &[4, 5]).unwrap();
    assert_eq!(2, buf.len());
    assert_eq!(4, buf.consumed_len());

    buf.reset();
    assert_eq!(0, buf.len());
    assert_eq!(0, buf.consumed_len());
}

#[test]
fn write_data_till_end_of_varint() {
    let mut buffer = new_receive_buffer();

    for nr_bytes in 0..64 * 2 + 1 {
        let data = vec![0u8; nr_bytes];
        assert!(buffer
            .write_at(
                VarInt::new(MAX_VARINT_VALUE - nr_bytes as u64).unwrap(),
                &data[..]
            )
            .is_ok());
    }
}

#[test]
fn fail_to_push_out_of_bounds_data() {
    let mut buffer = new_receive_buffer();

    for nr_bytes in 0..64 * 2 + 1 {
        let data = vec![0u8; nr_bytes + 1];
        assert_eq!(
            Err(ReceiveBufferError::OutOfRange),
            buffer.write_at(
                VarInt::new(MAX_VARINT_VALUE - nr_bytes as u64).unwrap(),
                &data[..]
            )
        );
    }
}

#[test]
#[cfg_attr(miri, ignore)] // miri fails because the slice points to invalid memory
#[cfg(target_pointer_width = "64")]
fn fail_to_push_out_of_bounds_data_with_long_buffer() {
    let mut buffer = ReceiveBuffer::new();

    // Overflow the allowed buffers by size 1. This uses an invalid memory
    // reference, due to not wanting to allocate too much memory. This is
    // ok in order to make sure that the overflow check works.
    // fake_data is based on a valid base pointer, in order to give us a
    // bit more certainty that the test won't fail for other reasons.
    let data = [0u8; 32];
    let fake_data = unsafe {
        core::slice::from_raw_parts(&data as *const u8, MAX_VARINT_VALUE as usize - 20 + 1)
    };

    for _ in 0..64 * 2 + 1 {
        assert_eq!(
            Err(ReceiveBufferError::OutOfRange),
            buffer.write_at(20u32.into(), fake_data)
        );
    }
}

#[test]
fn pop_watermarked_test() {
    let mut buffer = ReceiveBuffer::new();

    assert_eq!(
        None,
        buffer.pop_watermarked(1),
        "an empty buffer should always return none"
    );

    buffer
        .write_at(0u32.into(), &[0, 1, 2, 3, 4, 5, 6, 7])
        .unwrap();

    assert_eq!(
        None,
        buffer.pop_watermarked(0),
        "A 0 sized watermark should not pop the buffer"
    );

    assert_eq!(
        Some(bytes::BytesMut::from(&[0, 1, 2][..])),
        buffer.pop_watermarked(3),
        "the watermark should split off a piece of the buffer",
    );

    assert_eq!(
        Some(bytes::BytesMut::from(&[3][..])),
        buffer.pop_watermarked(1),
        "the watermark should split off another piece of the buffer",
    );

    assert_eq!(
        Some(bytes::BytesMut::from(&[4, 5, 6, 7][..])),
        buffer.pop_watermarked(100),
        "popping with a high watermark should remove the remaining chunk",
    );

    assert_eq!(
        None,
        buffer.pop_watermarked(1),
        "the receive buffer should be empty after splitting"
    );
}

#[test]
fn allocation_size_test() {
    // | bytes received | allocation size |
    // |----------------|-----------------|
    // | 0              | 4096            |
    // | 65536          | 16384           |
    // | 262144         | 32768           |
    // | 1048576        | 65536           |

    let received = [(0, 4096), (65536, 16384), (262144, 32768), (1048576, 65536)];

    for (index, (offset, size)) in received.iter().copied().enumerate() {
        assert_eq!(
            ReceiveBuffer::allocation_size(offset),
            size,
            "offset = {}",
            offset
        );

        if let Some((offset, _)) = received.get(index + 1) {
            let offset = offset - 1;
            assert_eq!(
                ReceiveBuffer::allocation_size(offset),
                size,
                "offset = {}",
                offset
            );
        }
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{application, crypto, transport};
pub use crate::{frame::ConnectionClose, inet::SocketAddress};

/// Provides a hook for applications to rewrite CONNECTION_CLOSE frames
///
/// Implementations should take care to not leak potentially sensitive information
/// to peers. This includes removing `reason` fields and making error codes more general.
pub trait Formatter: 'static + Send {
    /// Formats a transport error for use in 1-RTT (application data) packets
    fn format_transport_error(&self, context: &Context, error: transport::Error)
        -> ConnectionClose;

    /// Formats an application error for use in 1-RTT (application data) packets
    fn format_application_error(
        &self,
        context: &Context,
        error: application::Error,
    ) -> ConnectionClose;

    /// Formats a transport error for use in early (initial, handshake) packets
    fn format_early_transport_error(
        &self,
        context: &Context,
        error: transport::Error,
    ) -> ConnectionClose;

    /// Formats an application error for use in early (initial, handshake) packets
    fn format_early_application_error(
        &self,
        context: &Context,
        error: application::Error,
    ) -> ConnectionClose;
}

#[non_exhaustive]
#[derive(Debug)]
pub struct Context<'a> {
    pub remote_address: &'a SocketAddress,
}

impl<'a> Context<'a> {
    pub fn new(remote_address: &'a SocketAddress) -> Self {
        Self { remote_address }
    }
}

/// A formatter that passes errors through, unmodified
///
/// WARNING: This formatter should only be used in application development,
///          as it can leak potentially sensitive information to the peer.
#[derive(Clone, Copy, Debug, Default)]
pub struct Development;

impl Formatter for Development {
    fn format_transport_error(
        &self,
        _context: &Context,
        error: transport::Error,
    ) -> ConnectionClose {
        error.into()
    }

    fn format_application_error(
        &self,
        _context: &Context,
        error: application::Error,
    ) -> ConnectionClose {
        error.into()
    }

    fn format_early_transport_error(
        &self,
        _context: &Context,
        error: transport::Error,
    ) -> ConnectionClose {
        error.into()
    }

    fn format_early_application_error(
        &self,
        _context: &Context,
        error: application::Error,
    ) -> ConnectionClose {
        error.into()
    }
}

/// A formatter that removes potentially sensitive information
///
/// The following is performed:
///
/// * Reasons and frame_types are hidden
/// * INTERNAL_ERROR is transformed into PROTOCOL_VIOLATION
/// * Application codes are hidden in early (initial, handshake) packets
/// * Crypto (TLS) alerts are transformed into HANDSHAKE_FAILURE
#[derive(Clone, Copy, Debug, Default)]
pub struct Production;

impl Formatter for Production {
    fn format_transport_error(
        &self,
        _context: &Context,
        error: transport::Error,
    ) -> ConnectionClose {
        // rewrite internal errors as PROTOCOL_VIOLATION
        if error.code == transport::Error::INTERNAL_ERROR.code {
            return transport::Error::PROTOCOL_VIOLATION.into();
        }

        //= https://www.rfc-editor.org/rfc/rfc9001#section-4.8
        //# QUIC permits the use of a generic code in place of a specific error
        //# code; see Section 11 of [QUIC-TRANSPORT].  For TLS alerts, this
        //# includes replacing any alert with a generic alert, such as
        //# handshake_failure (0x0128 in QUIC).  Endpoints MAY use a generic
        //# error code to avoid possibly exposing confidential information.
        if error.try_into_crypto_error().is_some() {
            return transport::Error::from(crypto::CryptoError::HANDSHAKE_FAILURE).into();
        }

        // only preserve the error code
        transport::Error::new(error.code.as_varint()).into()
    }

    fn format_application_error(
        &self,
        _context: &Context,
        error: application::Error,
    ) -> ConnectionClose {
        error.into()
    }

    fn format_early_transport_error(
        &self,
        context: &Context,
        error: transport::Error,
    ) -> ConnectionClose {
        Self.format_transport_error(context, error)
    }

    fn format_early_application_error(
        &self,
        _context: &Context,
        _error: application::Error,
    ) -> ConnectionClose {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-10.2.3
        //# Sending a CONNECTION_CLOSE of type 0x1d in an Initial or Handshake
        //# packet could expose application state or be used to alter application
        //# state.  A CONNECTION_CLOSE of type 0x1d MUST be replaced by a
        //# CONNECTION_CLOSE of type 0x1c when sending the frame in Initial or
        //# Handshake packets.  Otherwise, information about the application
        //# state might be revealed.  Endpoints MUST clear the value of the
        //# Reason Phrase field and SHOULD use the APPLICATION_ERROR code when
        //# converting to a CONNECTION_CLOSE of type 0x1c.

        transport::Error::APPLICATION_ERROR.into()
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{
    application, connection, crypto::CryptoError, endpoint, frame::ConnectionClose, transport,
};
use core::{convert::TryInto, fmt, panic, time::Duration};

/// Errors that a connection can encounter.
#[derive(PartialEq, Eq, Debug, Copy, Clone)]
#[non_exhaustive]
pub enum Error {
    /// The connection was closed without an error
    #[non_exhaustive]
    Closed {
        initiator: endpoint::Location,
        source: &'static panic::Location<'static>,
    },

    /// The connection was closed on the transport level
    ///
    /// This can occur either locally or by the peer. The argument contains
    /// the error code which the transport provided in order to close the
    /// connection.
    #[non_exhaustive]
    Transport {
        code: transport::error::Code,
        frame_type: u64,
        reason: &'static str,
        initiator: endpoint::Location,
        source: &'static panic::Location<'static>,
    },

    /// The connection was closed on the application level
    ///
    /// This can occur either locally or by the peer. The argument contains
    /// the error code which the application/ supplied in order to close the
    /// connection.
    #[non_exhaustive]
    Application {
        error: application::Error,
        initiator: endpoint::Location,
        source: &'static panic::Location<'static>,
    },

    /// The connection was reset by a stateless reset from the peer
    #[non_exhaustive]
    StatelessReset {
        source: &'static panic::Location<'static>,
    },

    /// The connection was closed because the local connection's idle timer expired
    #[non_exhaustive]
    IdleTimerExpired {
        source: &'static panic::Location<'static>,
    },

    /// The connection was closed because there are no valid paths
    #[non_exhaustive]
    NoValidPath {
        source: &'static panic::Location<'static>,
    },

    /// All Stream IDs for Streams on the given connection had been exhausted
    #[non_exhaustive]
    StreamIdExhausted {
        source: &'static panic::Location<'static>,
    },

    /// The handshake has taken longer to complete than the configured max handshake duration
    #[non_exhaustive]
    MaxHandshakeDurationExceeded {
        max_handshake_duration: Duration,
        source: &'static panic::Location<'static>,
    },

    /// The connection should be closed immediately without notifying the peer
    #[non_exhaustive]
    ImmediateClose {
        reason: &'static str,
        source: &'static panic::Location<'static>,
    },

    /// The connection attempt was rejected because the endpoint is closing
    #[non_exhaustive]
    EndpointClosing {
        source: &'static panic::Location<'static>,
    },

    /// The connection was closed due to an unspecified reason
    #[non_exhaustive]
    Unspecified {
        source: &'static panic::Location<'static>,
    },
}

#[cfg(feature = "std")]
impl std::error::Error for Error {}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Self::Closed { initiator, .. } => write!(
                f,
                "The connection was closed without an error by {initiator}"
            ),
            Self::Transport { code, frame_type, reason, initiator, .. } => {
                let error = transport::Error {
                    code: *code,
                    frame_type: (*frame_type).try_into().ok().unwrap_or_default(),
                    reason,
                };
                write!(
                    f,
                    "The connection was closed on the transport level with error {error} by {initiator}"
                )
            },
            Self::Application { error, initiator, .. } => write!(
                f,
                "The connection was closed on the application level with error {error:?} by {initiator}"
            ),
            Self::StatelessReset { .. } => write!(
                f,
                "The connection was reset by a stateless reset by {}",
                endpoint::Location::Remote
            ),
            Self::IdleTimerExpired {.. } => write!(
                f,
                "The connection was closed because the connection's idle timer expired by {}",
                endpoint::Location::Local
            ),
            Self::NoValidPath { .. } => write!(
                f,
                "The connection was closed because there are no valid paths"
            ),
            Self::StreamIdExhausted { .. } => write!(
                f,
                "All Stream IDs for Streams on the given connection had been exhausted"
            ),
            Self::MaxHandshakeDurationExceeded { max_handshake_duration, .. } => write!(
              f,
                "The connection was closed because the handshake took longer than the max handshake \
                duration of {max_handshake_duration:?}"
            ),
            Self::ImmediateClose { reason, .. } => write!(
                f,
                "The connection was closed due to: {reason}"
            ),
            Self::EndpointClosing { .. } => {
                write!(f, "The connection attempt was rejected because the endpoint is closing")
            }
            Self::Unspecified { .. } => {
                write!(f, "The connection was closed due to an unspecified reason")
            }
        }
    }
}

impl Error {
    /// Returns the [`panic::Location`] for the error
    pub fn source(&self) -> &'static panic::Location<'static> {
        match self {
            Error::Closed { source, .. } => source,
            Error::Transport { source, .. } => source,
            Error::Application { source, .. } => source,
            Error::StatelessReset { source } => source,
            Error::IdleTimerExpired { source } => source,
            Error::NoValidPath { source } => source,
            Error::StreamIdExhausted { source } => source,
            Error::MaxHandshakeDurationExceeded { source, .. } => source,
            Error::ImmediateClose { source, .. } => source,
            Error::EndpointClosing { source } => source,
            Error::Unspecified { source } => source,
        }
    }

    #[track_caller]
    fn from_transport_error(error: transport::Error, initiator: endpoint::Location) -> Self {
        let source = panic::Location::caller();
        match error.code {
            // The connection closed without an error
            code if code == transport::Error::NO_ERROR.code => Self::Closed { initiator, source },
            // The connection closed without an error at the application layer
            code if code == transport::Error::APPLICATION_ERROR.code && initiator.is_remote() => {
                Self::Closed { initiator, source }
            }
            // The connection closed with an actual error
            _ => Self::Transport {
                code: error.code,
                frame_type: error.frame_type.into(),
                reason: error.reason,
                initiator,
                source,
            },
        }
    }

    #[inline]
    #[track_caller]
    #[doc(hidden)]
    pub fn closed(initiator: endpoint::Location) -> Error {
        let source = panic::Location::caller();
        Error::Closed { initiator, source }
    }

    #[inline]
    #[track_caller]
    #[doc(hidden)]
    pub fn immediate_close(reason: &'static str) -> Error {
        let source = panic::Location::caller();
        Error::ImmediateClose { reason, source }
    }

    #[inline]
    #[track_caller]
    #[doc(hidden)]
    pub fn idle_timer_expired() -> Error {
        let source = panic::Location::caller();
        Error::IdleTimerExpired { source }
    }

    #[inline]
    #[track_caller]
    #[doc(hidden)]
    pub fn stream_id_exhausted() -> Error {
        let source = panic::Location::caller();
        Error::StreamIdExhausted { source }
    }

    #[inline]
    #[track_caller]
    #[doc(hidden)]
    pub fn no_valid_path() -> Error {
        let source = panic::Location::caller();
        Error::NoValidPath { source }
    }

    #[inline]
    #[track_caller]
    #[doc(hidden)]
    pub fn stateless_reset() -> Error {
        let source = panic::Location::caller();
        Error::StatelessReset { source }
    }

    #[inline]
    #[track_caller]
    #[doc(hidden)]
    pub fn max_handshake_duration_exceeded(max_handshake_duration: Duration) -> Error {
        let source = panic::Location::caller();
        Error::MaxHandshakeDurationExceeded {
            max_handshake_duration,
            source,
        }
    }

    #[inline]
    #[track_caller]
    #[doc(hidden)]
    pub fn application(error: application::Error) -> Error {
        let source = panic::Location::caller();
        Error::Application {
            error,
            initiator: endpoint::Location::Local,
            source,
        }
    }

    #[inline]
    #[track_caller]
    #[doc(hidden)]
    pub fn endpoint_closing() -> Error {
        let source = panic::Location::caller();
        Error::EndpointClosing { source }
    }

    #[inline]
    #[track_caller]
    #[doc(hidden)]
    pub fn unspecified() -> Error {
        let source = panic::Location::caller();
        Error::Unspecified { source }
    }

    #[inline]
    #[doc(hidden)]
    pub fn into_accept_error(error: connection::Error) -> Result<(), connection::Error> {
        match error {
            // The connection closed without an error
            connection::Error::Closed { .. } => Ok(()),
            // The application closed the connection
            connection::Error::Transport { code, .. }
                if code == transport::Error::APPLICATION_ERROR.code =>
            {
                Ok(())
            }
            // The local connection's idle timer expired
            connection::Error::IdleTimerExpired { .. } => Ok(()),
            // Otherwise return the real error to the user
            _ => Err(error),
        }
    }
}

/// Returns a CONNECTION_CLOSE frame for the given connection Error, if any
///
/// The first item will be a close frame for an early (initial, handshake) packet.
/// The second item will be a close frame for a 1-RTT (application data) packet.
pub fn as_frame<'a, F: connection::close::Formatter>(
    error: Error,
    formatter: &'a F,
    context: &'a connection::close::Context<'a>,
) -> Option<(ConnectionClose<'a>, ConnectionClose<'a>)> {
    match error {
        Error::Closed { initiator, .. } => {
            // don't send CONNECTION_CLOSE frames on remote-initiated errors
            if initiator.is_remote() {
                return None;
            }

            let error = transport::Error::NO_ERROR;
            let early = formatter.format_early_transport_error(context, error);
            let one_rtt = formatter.format_transport_error(context, error);

            Some((early, one_rtt))
        }
        Error::Transport {
            code,
            frame_type,
            reason,
            initiator,
            ..
        } => {
            // don't send CONNECTION_CLOSE frames on remote-initiated errors
            if initiator.is_remote() {
                return None;
            }

            let error = transport::Error {
                code,
                frame_type: frame_type.try_into().unwrap_or_default(),
                reason,
            };

            let early = formatter.format_early_transport_error(context, error);
            let one_rtt = formatter.format_transport_error(context, error);
            Some((early, one_rtt))
        }
        Error::Application {
            error, initiator, ..
        } => {
            // don't send CONNECTION_CLOSE frames on remote-initiated errors
            if initiator.is_remote() {
                return None;
            }

            let early = formatter.format_early_application_error(context, error);
            let one_rtt = formatter.format_application_error(context, error);
            Some((early, one_rtt))
        }
        // This error comes from the peer so we don't respond with a CONNECTION_CLOSE
        Error::StatelessReset { .. } => None,
        // Nothing gets sent on idle timeouts
        Error::IdleTimerExpired { .. } => None,
        Error::NoValidPath { .. } => None,
        Error::StreamIdExhausted { .. } => {
            let error =
                transport::Error::PROTOCOL_VIOLATION.with_reason("stream IDs have been exhausted");

            let early = formatter.format_early_transport_error(context, error);
            let one_rtt = formatter.format_transport_error(context, error);

            Some((early, one_rtt))
        }
        Error::MaxHandshakeDurationExceeded { .. } => None,
        Error::ImmediateClose { .. } => None,
        Error::EndpointClosing { .. } => None,
        Error::Unspecified { .. } => {
            let error =
                transport::Error::INTERNAL_ERROR.with_reason("an unspecified error occurred");

            let early = formatter.format_early_transport_error(context, error);
            let one_rtt = formatter.format_transport_error(context, error);

            Some((early, one_rtt))
        }
    }
}

impl application::error::TryInto for Error {
    fn application_error(&self) -> Option<application::Error> {
        if let Self::Application { error, .. } = self {
            Some(*error)
        } else {
            None
        }
    }
}

impl From<transport::Error> for Error {
    #[track_caller]
    fn from(error: transport::Error) -> Self {
        Self::from_transport_error(error, endpoint::Location::Local)
    }
}

impl From<CryptoError> for Error {
    #[track_caller]
    fn from(error: CryptoError) -> Self {
        transport::Error::from(error).into()
    }
}

impl<'a> From<ConnectionClose<'a>> for Error {
    #[track_caller]
    fn from(error: ConnectionClose) -> Self {
        if let Some(frame_type) = error.frame_type {
            let error = transport::Error {
                code: transport::error::Code::new(error.error_code),
                // we use an empty `&'static str` so we don't allocate anything
                // in the event of an error
                reason: "",
                frame_type,
            };
            Self::from_transport_error(error, endpoint::Location::Remote)
        } else {
            let source = panic::Location::caller();
            Self::Application {
                error: error.error_code.into(),
                initiator: endpoint::Location::Remote,
                source,
            }
        }
    }
}

#[cfg(feature = "std")]
impl From<Error> for std::io::Error {
    fn from(error: Error) -> Self {
        let kind = error.into();
        std::io::Error::new(kind, error)
    }
}

#[cfg(feature = "std")]
impl From<Error> for std::io::ErrorKind {
    fn from(error: Error) -> Self {
        use std::io::ErrorKind;
        match error {
            Error::Closed { .. } => ErrorKind::ConnectionAborted,
            Error::Transport { code, .. } if code == transport::Error::CONNECTION_REFUSED.code => {
                ErrorKind::ConnectionRefused
            }
            Error::Transport { .. } => ErrorKind::ConnectionReset,
            Error::Application { .. } => ErrorKind::ConnectionReset,
            Error::StatelessReset { .. } => ErrorKind::ConnectionReset,
            Error::IdleTimerExpired { .. } => ErrorKind::TimedOut,
            Error::NoValidPath { .. } => ErrorKind::Other,
            Error::StreamIdExhausted { .. } => ErrorKind::Other,
            Error::MaxHandshakeDurationExceeded { .. } => ErrorKind::TimedOut,
            Error::ImmediateClose { .. } => ErrorKind::Other,
            Error::EndpointClosing { .. } => ErrorKind::Other,
            Error::Unspecified { .. } => ErrorKind::Other,
        }
    }
}

/// Some connection methods may need to indicate both `TransportError`s and `CryptoError`s. This
/// enum is used to allow for either error type to be returned as appropriate.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum ProcessingError {
    DuplicatePacket,
    /// Received a Retry packet with SCID field equal to DCID field.
    RetryScidEqualsDcid,
    ConnectionError(Error),
    CryptoError(CryptoError),
    NonEmptyRetryToken,
}

impl From<Error> for ProcessingError {
    fn from(inner_error: Error) -> Self {
        ProcessingError::ConnectionError(inner_error)
    }
}

impl From<crate::transport::Error> for ProcessingError {
    #[track_caller]
    fn from(inner_error: crate::transport::Error) -> Self {
        // Try extracting out the crypto error from other transport errors
        if let Some(error) = inner_error.try_into_crypto_error() {
            Self::CryptoError(error)
        } else {
            Self::ConnectionError(inner_error.into())
        }
    }
}

impl From<CryptoError> for ProcessingError {
    fn from(inner_error: CryptoError) -> Self {
        ProcessingError::CryptoError(inner_error)
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

//! Defines the QUIC connection ID

use crate::{
    event::{api::SocketAddress, IntoEvent},
    inet, transport,
};
use core::{
    convert::{TryFrom, TryInto},
    time::Duration,
};
use s2n_codec::{decoder_value, Encoder, EncoderValue};

#[cfg(any(test, feature = "generator"))]
use bolero_generator::*;

//= https://www.rfc-editor.org/rfc/rfc9000#section-5.1
//# Each connection possesses a set of connection identifiers, or
//# connection IDs, each of which can identify the connection.
//# Connection IDs are independently selected by endpoints; each endpoint
//# selects the connection IDs that its peer uses.

/// The maximum size of a connection ID.
pub const MAX_LEN: usize = crate::packet::long::DESTINATION_CONNECTION_ID_MAX_LEN;

/// The minimum lifetime of a connection ID.
pub const MIN_LIFETIME: Duration = Duration::from_secs(60);

/// The maximum bounded lifetime of a connection ID. Connection IDs may have no specified
/// lifetime at all, but if a lifetime is specified, it cannot exceed this value.
pub const MAX_LIFETIME: Duration = Duration::from_secs(24 * 60 * 60); // one day

macro_rules! id {
    ($type:ident, $min_len:expr) => {
        /// Uniquely identifies a QUIC connection between 2 peers
        #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
        #[cfg_attr(any(feature = "generator", test), derive(TypeGenerator))]
        pub struct $type {
            bytes: [u8; MAX_LEN],
            #[cfg_attr(any(feature = "generator", test), generator(Self::GENERATOR))]
            len: u8,
        }

        impl core::fmt::Debug for $type {
            fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
                write!(f, "{}({:?})", stringify!($type), self.as_bytes())
            }
        }

        impl $type {
            /// The minimum length for this connection ID type
            pub const MIN_LEN: usize = $min_len;

            #[cfg(any(feature = "generator", test))]
            const GENERATOR: core::ops::RangeInclusive<u8> = $min_len..=(MAX_LEN as u8);

            /// Creates a connection ID from a byte array.
            ///
            /// If the passed byte array exceeds the maximum allowed length for
            /// Connection IDs (20 bytes in QUIC v1) `None` will be returned.
            /// All other input values are valid.
            #[inline]
            pub fn try_from_bytes(bytes: &[u8]) -> Option<$type> {
                Self::try_from(bytes).ok()
            }

            /// Returns the Connection ID in byte form
            #[inline]
            pub fn as_bytes(&self) -> &[u8] {
                self.as_ref()
            }

            /// Returns the length of the connection id
            #[inline]
            pub const fn len(&self) -> usize {
                self.len as usize
            }

            /// Returns true if this connection ID is zero-length
            #[inline]
            pub fn is_empty(&self) -> bool {
                self.len == 0
            }

            /// A connection ID to use for testing
            #[cfg(any(test, feature = "testing"))]
            pub const TEST_ID: Self = Self::test_id();

            // Constructs a test connection ID by converting the name
            // of the ID type to bytes and populating the first and last
            // bytes of a max length ID with those bytes.
            #[cfg(any(test, feature = "testing"))]
            const fn test_id() -> Self {
                let type_bytes = stringify!($type).as_bytes();
                let mut result = [0u8; MAX_LEN];
                result[0] = type_bytes[0];
                result[1] = type_bytes[1];
                result[2] = type_bytes[2];
                result[3] = type_bytes[3];
                result[4] = type_bytes[4];
                result[5] = type_bytes[5];
                result[14] = type_bytes[0];
                result[15] = type_bytes[1];
                result[16] = type_bytes[2];
                result[17] = type_bytes[3];
                result[18] = type_bytes[4];
                result[19] = type_bytes[5];
                Self {
                    bytes: result,
                    len: MAX_LEN as u8,
                }
            }
        }

        impl From<[u8; MAX_LEN]> for $type {
            #[inline]
            fn from(bytes: [u8; MAX_LEN]) -> Self {
                Self {
                    bytes,
                    len: MAX_LEN as u8,
                }
            }
        }

        impl TryFrom<&[u8]> for $type {
            type Error = Error;

            #[inline]
            fn try_from(slice: &[u8]) -> Result<Self, Self::Error> {
                let len = slice.len();
                if !($type::MIN_LEN..=MAX_LEN).contains(&len) {
                    return Err(Error::InvalidLength);
                }
                let mut bytes = [0; MAX_LEN];
                bytes[..len].copy_from_slice(slice);
                Ok(Self {
                    bytes,
                    len: len as u8,
                })
            }
        }

        impl AsRef<[u8]> for $type {
            #[inline]
            fn as_ref(&self) -> &[u8] {
                &self.bytes[0..self.len as usize]
            }
        }

        decoder_value!(
            impl<'a> $type {
                fn decode(buffer: Buffer) -> Result<Self> {
                    let len = buffer.len();
                    let (value, buffer) = buffer.decode_slice(len)?;
                    let value: &[u8] = value.into_less_safe_slice();
                    let connection_id = $type::try_from(value).map_err(|_| {
                        s2n_codec::DecoderError::InvariantViolation(concat!(
                            "invalid ",
                            stringify!($type)
                        ))
                    })?;

                    Ok((connection_id, buffer))
                }
            }
        );

        impl EncoderValue for $type {
            #[inline]
            fn encode<E: Encoder>(&self, encoder: &mut E) {
                self.as_ref().encode(encoder)
            }
        }

        // Implement Default to allow for transport_parameter macro to work consistently,
        // though this value should never be used.
        impl Default for $type {
            fn default() -> Self {
                unimplemented!("connection IDs do not have default values")
            }
        }
    };
}

// Connection IDs that are generated locally and used to route packets from the peer to the local
// endpoint. s2n-QUIC does not provide zero-length connection IDs, the minimum allowable LocalId
// is 4 bytes.
id!(LocalId, 4);

// Connection IDs used to route packets to the peer. The peer may choose to use zero-length
// connection IDs.
id!(PeerId, 0);

// Connection IDs that are used as either a LocalId or a PeerId depending on if the endpoint is a
// server or a client, and thus the minimum length of the ID is not validated.
id!(UnboundedId, 0);
// The randomly generated ID the client sends when first contacting a server.
//= https://www.rfc-editor.org/rfc/rfc9000#section-7.2
//# When an Initial packet is sent by a client that has not previously
//# received an Initial or Retry packet from the server, the client
//# populates the Destination Connection ID field with an unpredictable
//# value.  This Destination Connection ID MUST be at least 8 bytes in
//# length.
id!(InitialId, 8);

impl From<LocalId> for UnboundedId {
    fn from(id: LocalId) -> Self {
        UnboundedId {
            bytes: id.bytes,
            len: id.len,
        }
    }
}

impl From<PeerId> for UnboundedId {
    fn from(id: PeerId) -> Self {
        UnboundedId {
            bytes: id.bytes,
            len: id.len,
        }
    }
}

impl From<InitialId> for UnboundedId {
    fn from(id: InitialId) -> Self {
        UnboundedId {
            bytes: id.bytes,
            len: id.len,
        }
    }
}

impl From<InitialId> for PeerId {
    fn from(id: InitialId) -> Self {
        PeerId {
            bytes: id.bytes,
            len: id.len,
        }
    }
}

// A LocalId may be converted to an InitialId, but InitialId has a higher
// minimum length, so conversion may not succeed.
impl TryFrom<LocalId> for InitialId {
    type Error = Error;

    #[inline]
    fn try_from(value: LocalId) -> Result<Self, Self::Error> {
        value.as_bytes().try_into()
    }
}

#[derive(Debug, PartialEq)]
pub enum Error {
    InvalidLength,
    InvalidLifetime,
}

impl Error {
    fn message(&self) -> &'static str {
        match self {
            Error::InvalidLength => "invalid connection id length",
            Error::InvalidLifetime => "invalid connection id lifetime",
        }
    }
}

impl core::fmt::Display for Error {
    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
        write!(f, "{}", self.message())
    }
}

impl From<Error> for transport::Error {
    #[inline]
    fn from(error: Error) -> Self {
        Self::PROTOCOL_VIOLATION.with_reason(error.message())
    }
}

/// Information about the connection that may be used
/// when generating or validating connection IDs
#[derive(Clone, Debug)]
#[non_exhaustive]
pub struct ConnectionInfo<'a> {
    pub remote_address: SocketAddress<'a>,
}

impl<'a> ConnectionInfo<'a> {
    #[inline]
    #[doc(hidden)]
    pub fn new(remote_address: &'a inet::SocketAddress) -> Self {
        Self {
            remote_address: remote_address.into_event(),
        }
    }
}

/// Format for connection IDs
pub trait Format: 'static + Validator + Generator + Send {}

/// Implement Format for all types that implement the required subtraits
impl<T: 'static + Validator + Generator + Send> Format for T {}

/// A validator for a connection ID format
pub trait Validator {
    //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3.2
    //# An endpoint that uses this design MUST
    //# either use the same connection ID length for all connections or
    //# encode the length of the connection ID such that it can be recovered
    //# without state.
    /// Validates a connection ID from a buffer
    ///
    /// Implementations should handle situations where the buffer will include extra
    /// data after the connection ID.
    ///
    /// Returns the length of the connection id if successful, otherwise `None` is returned.
    fn validate(&self, connection_info: &ConnectionInfo, buffer: &[u8]) -> Option<usize>;
}

impl Validator for usize {
    #[inline]
    fn validate(&self, _connection_info: &ConnectionInfo, buffer: &[u8]) -> Option<usize> {
        if buffer.len() >= *self {
            Some(*self)
        } else {
            None
        }
    }
}

/// A generator for a connection ID format
pub trait Generator {
    /// Generates a connection ID.
    ///
    /// Connection IDs MUST NOT contain any information that can be used by
    /// an external observer (that is, one that does not cooperate with the
    /// issuer) to correlate them with other connection IDs for the same
    /// connection.
    ///
    /// Each call to `generate` should produce a unique Connection ID,
    /// otherwise the endpoint may terminate.
    fn generate(&mut self, connection_info: &ConnectionInfo) -> LocalId;

    /// The maximum amount of time each generated connection ID should be
    /// used for. By default there is no maximum, though connection IDs
    /// may be retired due to rotation requirements or peer requests.
    #[inline]
    fn lifetime(&self) -> Option<core::time::Duration> {
        None
    }
}

#[derive(Debug, Copy, Clone, Default, PartialEq, Eq, PartialOrd, Ord)]
pub enum Interest {
    /// No new connection Ids are required
    #[default]
    None,
    /// The specified number of new connection Ids are required
    New(u8),
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn create_connection_id() {
        let connection_id = LocalId::try_from_bytes(b"My Connection 123").unwrap();
        assert_eq!(b"My Connection 123", connection_id.as_bytes());

        let connection_id = PeerId::try_from_bytes(b"My Connection 456").unwrap();
        assert_eq!(b"My Connection 456", connection_id.as_bytes());

        let connection_id = InitialId::try_from_bytes(b"My Connection 789").unwrap();
        assert_eq!(b"My Connection 789", connection_id.as_bytes());
    }

    #[test]
    fn exceed_max_connection_id_length() {
        let connection_id_bytes = [0u8; MAX_LEN];
        assert!(LocalId::try_from_bytes(&connection_id_bytes).is_some());
        assert!(PeerId::try_from_bytes(&connection_id_bytes).is_some());
        assert!(InitialId::try_from_bytes(&connection_id_bytes).is_some());

        let connection_id_bytes = [0u8; MAX_LEN + 1];
        assert!(LocalId::try_from_bytes(&connection_id_bytes).is_none());
        assert!(PeerId::try_from_bytes(&connection_id_bytes).is_none());
        assert!(InitialId::try_from_bytes(&connection_id_bytes).is_none());
    }

    #[test]
    fn min_connection_id_length() {
        let connection_id_bytes = [0u8; LocalId::MIN_LEN];
        assert!(LocalId::try_from_bytes(&connection_id_bytes).is_some());

        let connection_id_bytes = [0u8; PeerId::MIN_LEN];
        assert!(PeerId::try_from_bytes(&connection_id_bytes).is_some());

        let connection_id_bytes = [0u8; InitialId::MIN_LEN];
        assert!(InitialId::try_from_bytes(&connection_id_bytes).is_some());

        let connection_id_bytes = [0u8; LocalId::MIN_LEN - 1];
        assert!(LocalId::try_from_bytes(&connection_id_bytes).is_none());

        let connection_id_bytes = [0u8; InitialId::MIN_LEN - 1];
        assert!(InitialId::try_from_bytes(&connection_id_bytes).is_none());
    }

    #[test]
    fn unbounded_id() {
        let connection_id_bytes = [0u8; LocalId::MIN_LEN];
        assert!(UnboundedId::try_from_bytes(&connection_id_bytes).is_some());

        let connection_id_bytes = [0u8; PeerId::MIN_LEN];
        assert!(UnboundedId::try_from_bytes(&connection_id_bytes).is_some());

        let connection_id_bytes = [0u8; InitialId::MIN_LEN];
        assert!(UnboundedId::try_from_bytes(&connection_id_bytes).is_some());

        println!("{:?}", LocalId::TEST_ID);
        println!("{:?}", PeerId::TEST_ID);
        println!("{:?}", UnboundedId::TEST_ID);
    }
}

#[cfg(any(test, feature = "testing"))]
pub mod testing {
    use super::*;
    use core::convert::TryInto;

    #[derive(Debug, Default)]
    pub struct Format(u64);

    impl Validator for Format {
        fn validate(&self, _connection_info: &ConnectionInfo, _buffer: &[u8]) -> Option<usize> {
            Some(core::mem::size_of::<u64>())
        }
    }

    impl Generator for Format {
        fn generate(&mut self, _connection_info: &ConnectionInfo) -> LocalId {
            let id = (&self.0.to_be_bytes()[..]).try_into().unwrap();
            self.0 += 1;
            id
        }
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{
    ack,
    event::{api::SocketAddress, IntoEvent},
    inet, stream,
    transport::parameters::{
        AckDelayExponent, ActiveConnectionIdLimit, InitialFlowControlLimits, InitialMaxData,
        InitialMaxStreamDataBidiLocal, InitialMaxStreamDataBidiRemote, InitialMaxStreamDataUni,
        InitialMaxStreamsBidi, InitialMaxStreamsUni, InitialStreamLimits, MaxAckDelay,
        MaxDatagramFrameSize, MaxIdleTimeout, TransportParameters,
    },
};
use core::{convert::TryInto, time::Duration};
use s2n_codec::decoder_invariant;

pub use crate::transport::parameters::ValidationError;

const MAX_HANDSHAKE_DURATION_DEFAULT: Duration = Duration::from_secs(10);

//= https://www.rfc-editor.org/rfc/rfc9000#section-10.1.2
//# A connection will time out if no packets are sent or received for a
//# period longer than the time negotiated using the max_idle_timeout
//# transport parameter; see Section 10.  However, state in middleboxes
//# might time out earlier than that.  Though REQ-5 in [RFC4787]
//# recommends a 2-minute timeout interval, experience shows that sending
//# packets every 30 seconds is necessary to prevent the majority of
//# middleboxes from losing state for UDP flows [GATEWAY].
const MAX_KEEP_ALIVE_PERIOD_DEFAULT: Duration = Duration::from_secs(30);

#[non_exhaustive]
#[derive(Debug)]
pub struct ConnectionInfo<'a> {
    pub remote_address: SocketAddress<'a>,
}

impl<'a> ConnectionInfo<'a> {
    #[inline]
    #[doc(hidden)]
    pub fn new(remote_address: &'a inet::SocketAddress) -> Self {
        Self {
            remote_address: remote_address.into_event(),
        }
    }
}

#[derive(Clone, Copy, Debug)]
pub struct Limits {
    pub(crate) max_idle_timeout: MaxIdleTimeout,
    pub(crate) data_window: InitialMaxData,
    pub(crate) bidirectional_local_data_window: InitialMaxStreamDataBidiLocal,
    pub(crate) bidirectional_remote_data_window: InitialMaxStreamDataBidiRemote,
    pub(crate) unidirectional_data_window: InitialMaxStreamDataUni,
    pub(crate) max_open_local_bidirectional_streams: stream::limits::LocalBidirectional,
    pub(crate) max_open_local_unidirectional_streams: stream::limits::LocalUnidirectional,
    pub(crate) max_open_remote_bidirectional_streams: InitialMaxStreamsBidi,
    pub(crate) max_open_remote_unidirectional_streams: InitialMaxStreamsUni,
    pub(crate) max_ack_delay: MaxAckDelay,
    pub(crate) ack_delay_exponent: AckDelayExponent,
    pub(crate) max_active_connection_ids: ActiveConnectionIdLimit,
    pub(crate) ack_elicitation_interval: u8,
    pub(crate) ack_ranges_limit: u8,
    pub(crate) max_send_buffer_size: stream::limits::MaxSendBufferSize,
    pub(crate) max_handshake_duration: Duration,
    pub(crate) max_keep_alive_period: Duration,
    pub(crate) max_datagram_frame_size: MaxDatagramFrameSize,
}

impl Default for Limits {
    fn default() -> Self {
        Self::new()
    }
}

macro_rules! setter {
    ($name:ident, $field:ident, $inner:ty $(, |$validate_value:ident| $validaiton:block)?) => {
        pub fn $name(mut self, value: $inner) -> Result<Self, ValidationError> {
            $(
                let $validate_value = value;
                $validaiton
            )?
            self.$field = value.try_into()?;
            Ok(self)
        }
    };
}

impl Limits {
    pub const fn new() -> Self {
        Self {
            max_idle_timeout: MaxIdleTimeout::RECOMMENDED,
            data_window: InitialMaxData::RECOMMENDED,
            bidirectional_local_data_window: InitialMaxStreamDataBidiLocal::RECOMMENDED,
            bidirectional_remote_data_window: InitialMaxStreamDataBidiRemote::RECOMMENDED,
            unidirectional_data_window: InitialMaxStreamDataUni::RECOMMENDED,
            max_open_local_bidirectional_streams: stream::limits::LocalBidirectional::RECOMMENDED,
            max_open_local_unidirectional_streams: stream::limits::LocalUnidirectional::RECOMMENDED,
            max_open_remote_bidirectional_streams: InitialMaxStreamsBidi::RECOMMENDED,
            max_open_remote_unidirectional_streams: InitialMaxStreamsUni::RECOMMENDED,
            max_ack_delay: MaxAckDelay::RECOMMENDED,
            ack_delay_exponent: AckDelayExponent::RECOMMENDED,
            max_active_connection_ids: ActiveConnectionIdLimit::RECOMMENDED,
            ack_elicitation_interval: ack::Settings::RECOMMENDED.ack_elicitation_interval,
            ack_ranges_limit: ack::Settings::RECOMMENDED.ack_ranges_limit,
            max_send_buffer_size: stream::Limits::RECOMMENDED.max_send_buffer_size,
            max_handshake_duration: MAX_HANDSHAKE_DURATION_DEFAULT,
            max_keep_alive_period: MAX_KEEP_ALIVE_PERIOD_DEFAULT,
            max_datagram_frame_size: MaxDatagramFrameSize::DEFAULT,
        }
    }

    // We limit the initial data limit to u32::MAX (4GB), which far
    // exceeds the reasonable amount of data a connection is
    // initially allowed to send.
    //
    // By representing the flow control value as a u32, we save space
    // on the connection state.
    setter!(with_data_window, data_window, u64, |validate_value| {
        decoder_invariant!(
            validate_value <= u32::MAX.into(),
            "data_window must be <= u32::MAX"
        );
    });
    setter!(
        with_bidirectional_local_data_window,
        bidirectional_local_data_window,
        u64,
        |validate_value| {
            decoder_invariant!(
                validate_value <= u32::MAX.into(),
                "bidirectional_local_data_window must be <= u32::MAX"
            );
        }
    );
    setter!(
        with_bidirectional_remote_data_window,
        bidirectional_remote_data_window,
        u64,
        |validate_value| {
            decoder_invariant!(
                validate_value <= u32::MAX.into(),
                "bidirectional_remote_data_window must be <= u32::MAX"
            );
        }
    );
    setter!(
        with_unidirectional_data_window,
        unidirectional_data_window,
        u64,
        |validate_value| {
            decoder_invariant!(
                validate_value <= u32::MAX.into(),
                "unidirectional_data_window must be <= u32::MAX"
            );
        }
    );

    setter!(with_max_idle_timeout, max_idle_timeout, Duration);

    /// Sets both the max local and remote limits for bidirectional streams.
    #[deprecated(
        note = "use with_max_open_local_bidirectional_streams and with_max_open_remote_bidirectional_streams instead"
    )]
    pub fn with_max_open_bidirectional_streams(
        mut self,
        value: u64,
    ) -> Result<Self, ValidationError> {
        self.max_open_local_bidirectional_streams = value.try_into()?;
        self.max_open_remote_bidirectional_streams = value.try_into()?;
        Ok(self)
    }

    /// Sets the max local limits for bidirectional streams
    ///
    /// The value set is used instead of `with_max_open_bidirectional_streams` when set.
    pub fn with_max_open_local_bidirectional_streams(
        mut self,
        value: u64,
    ) -> Result<Self, ValidationError> {
        self.max_open_local_bidirectional_streams = value.try_into()?;
        Ok(self)
    }

    /// Sets the max remote limits for bidirectional streams.
    ///
    /// The value set is used instead of `with_max_open_bidirectional_streams` when set.
    pub fn with_max_open_remote_bidirectional_streams(
        mut self,
        value: u64,
    ) -> Result<Self, ValidationError> {
        self.max_open_remote_bidirectional_streams = value.try_into()?;
        Ok(self)
    }

    setter!(
        with_max_open_local_unidirectional_streams,
        max_open_local_unidirectional_streams,
        u64
    );
    setter!(
        with_max_open_remote_unidirectional_streams,
        max_open_remote_unidirectional_streams,
        u64
    );
    setter!(with_max_ack_delay, max_ack_delay, Duration);
    setter!(
        with_max_active_connection_ids,
        max_active_connection_ids,
        u64
    );
    setter!(with_ack_elicitation_interval, ack_elicitation_interval, u8);
    setter!(with_max_ack_ranges, ack_ranges_limit, u8);
    setter!(with_max_send_buffer_size, max_send_buffer_size, u32);
    setter!(
        with_max_handshake_duration,
        max_handshake_duration,
        Duration
    );
    setter!(with_max_keep_alive_period, max_keep_alive_period, Duration);

    // internal APIs

    #[doc(hidden)]
    #[inline]
    pub fn load_peer<A, B, C, D>(&mut self, peer_parameters: &TransportParameters<A, B, C, D>) {
        self.max_idle_timeout
            .load_peer(&peer_parameters.max_idle_timeout);
    }

    #[doc(hidden)]
    #[inline]
    pub const fn ack_settings(&self) -> ack::Settings {
        ack::Settings {
            ack_delay_exponent: self.ack_delay_exponent.as_u8(),
            max_ack_delay: self.max_ack_delay.as_duration(),
            ack_ranges_limit: self.ack_ranges_limit,
            ack_elicitation_interval: self.ack_elicitation_interval,
        }
    }

    #[doc(hidden)]
    #[inline]
    pub const fn initial_flow_control_limits(&self) -> InitialFlowControlLimits {
        InitialFlowControlLimits {
            stream_limits: self.initial_stream_limits(),
            max_data: self.data_window.as_varint(),
            max_open_remote_bidirectional_streams: self
                .max_open_remote_bidirectional_streams
                .as_varint(),
            max_open_remote_unidirectional_streams: self
                .max_open_remote_unidirectional_streams
                .as_varint(),
        }
    }

    #[doc(hidden)]
    #[inline]
    pub const fn initial_stream_limits(&self) -> InitialStreamLimits {
        InitialStreamLimits {
            max_data_bidi_local: self.bidirectional_local_data_window.as_varint(),
            max_data_bidi_remote: self.bidirectional_remote_data_window.as_varint(),
            max_data_uni: self.unidirectional_data_window.as_varint(),
        }
    }

    #[doc(hidden)]
    #[inline]
    pub fn stream_limits(&self) -> stream::Limits {
        stream::Limits {
            max_send_buffer_size: self.max_send_buffer_size,
            max_open_local_unidirectional_streams: self.max_open_local_unidirectional_streams,
            max_open_local_bidirectional_streams: self.max_open_local_bidirectional_streams,
        }
    }

    #[doc(hidden)]
    #[inline]
    pub fn max_idle_timeout(&self) -> Option<Duration> {
        self.max_idle_timeout.as_duration()
    }

    #[doc(hidden)]
    #[inline]
    pub fn max_handshake_duration(&self) -> Duration {
        self.max_handshake_duration
    }

    #[doc(hidden)]
    #[inline]
    pub fn max_keep_alive_period(&self) -> Duration {
        self.max_keep_alive_period
    }
}

/// Creates limits for a given connection
pub trait Limiter: 'static + Send {
    fn on_connection(&mut self, info: &ConnectionInfo) -> Limits;
}

/// Implement Limiter for a Limits struct
impl Limiter for Limits {
    fn on_connection(&mut self, _into: &ConnectionInfo) -> Limits {
        *self
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    // Local max data limits should be <= u32::MAX
    #[test]
    fn limit_validation() {
        let mut data = u32::MAX as u64 + 1;
        let limits = Limits::default();
        assert!(limits.with_data_window(data).is_err());
        assert!(limits.with_bidirectional_local_data_window(data).is_err());
        assert!(limits.with_bidirectional_remote_data_window(data).is_err());
        assert!(limits.with_unidirectional_data_window(data).is_err());

        data = u32::MAX as u64;
        assert!(limits.with_data_window(data).is_ok());
        assert!(limits.with_bidirectional_local_data_window(data).is_ok());
        assert!(limits.with_bidirectional_remote_data_window(data).is_ok());
        assert!(limits.with_unidirectional_data_window(data).is_ok());
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

pub mod close;
pub mod error;
pub mod id;
pub mod limits;

pub use error::{Error, ProcessingError};
pub use id::{InitialId, LocalId, PeerId, UnboundedId};
pub use limits::Limits;


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use core::fmt;
use s2n_codec::DecoderError;

/// Error type for crypto-related errors
#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[cfg_attr(feature = "thiserror", derive(thiserror::Error))]
pub struct CryptoError {
    pub reason: &'static str,
    pub code: u8,
}

impl CryptoError {
    /// Creates a new `CryptoError`
    pub const fn new(code: u8) -> Self {
        Self { code, reason: "" }
    }

    /// Sets the reason for `CryptoError`
    #[must_use]
    pub const fn with_reason(mut self, reason: &'static str) -> Self {
        self.reason = reason;
        self
    }
}

impl fmt::Display for CryptoError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if !self.reason.is_empty() {
            self.reason.fmt(f)
        } else if let Some(description) = self.description() {
            description.fmt(f)
        } else {
            write!(f, "CryptoError({})", self.code)
        }
    }
}

impl fmt::Debug for CryptoError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut d = f.debug_struct("CryptoError");

        d.field("code", &self.code);

        if let Some(description) = self.description() {
            d.field("description", &description);
        }

        if !self.reason.is_empty() {
            d.field("reason", &self.reason);
        }

        d.finish()
    }
}

impl From<DecoderError> for CryptoError {
    fn from(_: DecoderError) -> Self {
        Self::DECODE_ERROR
    }
}

macro_rules! alert_descriptions {
    ($($name:ident = $value:expr),* $(,)?) => {
        impl CryptoError {
            pub fn description(&self) -> Option<&'static str> {
                match self.code {
                    $(
                        $value => Some(stringify!($name)),
                    )*
                    _ => None,
                }
            }

            $(
                pub const $name: Self = Self::new($value);
            )*
        }

        #[test]
        fn description_test() {
            $(
                assert_eq!(&CryptoError::$name.to_string(), stringify!($name));
            )*
        }
    };
}

//= https://www.rfc-editor.org/rfc/rfc8446#appendix-B.2
//# enum { warning(1), fatal(2), (255) } AlertLevel;
//#
//# enum {
//#     close_notify(0),
//#     unexpected_message(10),
//#     bad_record_mac(20),
//#     decryption_failed_RESERVED(21),
//#     record_overflow(22),
//#     decompression_failure_RESERVED(30),
//#     handshake_failure(40),
//#     no_certificate_RESERVED(41),
//#     bad_certificate(42),
//#     unsupported_certificate(43),
//#     certificate_revoked(44),
//#     certificate_expired(45),
//#     certificate_unknown(46),
//#     illegal_parameter(47),
//#     unknown_ca(48),
//#     access_denied(49),
//#     decode_error(50),
//#     decrypt_error(51),
//#     export_restriction_RESERVED(60),
//#     protocol_version(70),
//#     insufficient_security(71),
//#     internal_error(80),
//#     inappropriate_fallback(86),
//#     user_canceled(90),
//#     no_renegotiation_RESERVED(100),
//#     missing_extension(109),
//#     unsupported_extension(110),
//#     certificate_unobtainable_RESERVED(111),
//#     unrecognized_name(112),
//#     bad_certificate_status_response(113),
//#     bad_certificate_hash_value_RESERVED(114),
//#     unknown_psk_identity(115),
//#     certificate_required(116),
//#     no_application_protocol(120),
//#     (255)
//# } AlertDescription;
//#
//# struct {
//#     AlertLevel level;
//#     AlertDescription description;
//# } Alert;

alert_descriptions!(
    CLOSE_NOTIFY = 0,
    UNEXPECTED_MESSAGE = 10,
    BAD_RECORD_MAC = 20,
    DECRYPTION_FAILED_RESERVED = 21,
    RECORD_OVERFLOW = 22,
    DECOMPRESSION_FAILURE_RESERVED = 30,
    HANDSHAKE_FAILURE = 40,
    NO_CERTIFICATE_RESERVED = 41,
    BAD_CERTIFICATE = 42,
    UNSUPPORTED_CERTIFICATE = 43,
    CERTIFICATE_REVOKED = 44,
    CERTIFICATE_EXPIRED = 45,
    CERTIFICATE_UNKNOWN = 46,
    ILLEGAL_PARAMETER = 47,
    UNKNOWN_CA = 48,
    ACCESS_DENIED = 49,
    DECODE_ERROR = 50,
    DECRYPT_ERROR = 51,
    EXPORT_RESTRICTION_RESERVED = 60,
    PROTOCOL_VERSION = 70,
    INSUFFICIENT_SECURITY = 71,
    INTERNAL_ERROR = 80,
    INAPPROPRIATE_FALLBACK = 86,
    USER_CANCELED = 90,
    NO_RENEGOTIATION_RESERVED = 100,
    MISSING_EXTENSION = 109,
    UNSUPPORTED_EXTENSION = 110,
    CERTIFICATE_UNOBTAINABLE_RESERVED = 111,
    UNRECOGNIZED_NAME = 112,
    BAD_CERTIFICATE_STATUS_RESPONSE = 113,
    BAD_CERTIFICATE_HASH_VALUE_RESERVED = 114,
    UNKNOWN_PSK_IDENTITY = 115,
    CERTIFICATE_REQUIRED = 116,
    NO_APPLICATION_PROTOCOL = 120,
);


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::crypto::{HeaderKey, Key};

/// Types for which are able to perform handshake cryptography.
///
/// This marker trait ensures only Handshake-level keys
/// are used with Handshake packets. Any key misuses are
/// caught by the type system.
pub trait HandshakeKey: Key {}

/// Types for which are able to perform handshake header cryptography.
///
/// This marker trait ensures only Handshake-level header keys
/// are used with Handshake packets. Any key misuses are
/// caught by the type system.
pub trait HandshakeHeaderKey: HeaderKey {}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{
    crypto::{EncryptedPayload, ProtectedPayload},
    packet::number::{PacketNumberSpace, TruncatedPacketNumber},
};
use s2n_codec::{DecoderBuffer, DecoderError};

/// Types for which are able to perform header cryptography.
pub trait HeaderKey: Send {
    /// Derives a header protection mask from a sample buffer, to be
    /// used for opening a packet.
    ///
    /// The sample size is determined by the key function.
    fn opening_header_protection_mask(&self, ciphertext_sample: &[u8]) -> HeaderProtectionMask;

    /// Returns the sample size needed for the header protection
    /// buffer
    fn opening_sample_len(&self) -> usize;

    /// Derives a header protection mask from a sample buffer, to be
    /// used for sealing a packet.
    ///
    /// The sample size is determined by the key function.
    fn sealing_header_protection_mask(&self, ciphertext_sample: &[u8]) -> HeaderProtectionMask;

    /// Returns the sample size needed for the header protection
    /// buffer
    fn sealing_sample_len(&self) -> usize;
}

//= https://www.rfc-editor.org/rfc/rfc9001#section-5.4.1
//# The output of this algorithm is a 5 byte mask that is applied to the
//# protected header fields using exclusive OR.

pub const HEADER_PROTECTION_MASK_LEN: usize = 5;
pub type HeaderProtectionMask = [u8; HEADER_PROTECTION_MASK_LEN];

//= https://www.rfc-editor.org/rfc/rfc9001#section-5.4.1
//# Figure 6 shows a sample algorithm for applying header protection.
//# Removing header protection only differs in the order in which the
//# packet number length (pn_length) is determined (here "^" is used to
//# represent exclusive OR).
//#
//# mask = header_protection(hp_key, sample)
//#
//# pn_length = (packet[0] & 0x03) + 1
//# if (packet[0] & 0x80) == 0x80:
//# # Long header: 4 bits masked
//# packet[0] ^= mask[0] & 0x0f
//# else:
//# # Short header: 5 bits masked
//# packet[0] ^= mask[0] & 0x1f
//#
//# # pn_offset is the start of the Packet Number field.
//# packet[pn_offset:pn_offset+pn_length] ^= mask[1:1+pn_length]

const LONG_HEADER_TAG: u8 = 0x80;
pub(crate) const LONG_HEADER_MASK: u8 = 0x0f;
pub(crate) const SHORT_HEADER_MASK: u8 = 0x1f;

#[inline(always)]
fn mask_from_packet_tag(tag: u8) -> u8 {
    if tag & LONG_HEADER_TAG == LONG_HEADER_TAG {
        LONG_HEADER_MASK
    } else {
        SHORT_HEADER_MASK
    }
}

#[inline(always)]
fn xor_mask(payload: &mut [u8], mask: &[u8]) {
    for (payload_byte, mask_byte) in payload.iter_mut().zip(&mask[1..]) {
        *payload_byte ^= mask_byte;
    }
}

#[inline]
pub(crate) fn apply_header_protection(
    mask: HeaderProtectionMask,
    payload: EncryptedPayload,
) -> ProtectedPayload {
    let header_len = payload.header_len;
    let packet_number_len = payload.packet_number_len;
    let payload = payload.buffer.into_less_safe_slice();

    payload[0] ^= mask[0] & mask_from_packet_tag(payload[0]);

    let header_with_pn_len = packet_number_len.bytesize() + header_len;
    let packet_number_bytes = &mut payload[header_len..header_with_pn_len];
    xor_mask(packet_number_bytes, &mask);

    ProtectedPayload::new(header_len, payload)
}

#[inline]
pub(crate) fn remove_header_protection(
    space: PacketNumberSpace,
    mask: HeaderProtectionMask,
    payload: ProtectedPayload,
) -> Result<(TruncatedPacketNumber, EncryptedPayload), DecoderError> {
    let header_len = payload.header_len;
    let payload = payload.buffer.into_less_safe_slice();

    payload[0] ^= mask[0] & mask_from_packet_tag(payload[0]);
    let packet_number_len = space.new_packet_number_len(payload[0]);

    let header_with_pn_len = packet_number_len.bytesize() + header_len;
    let packet_number = {
        let packet_number_bytes = &mut payload[header_len..header_with_pn_len];
        xor_mask(packet_number_bytes, &mask);

        let (packet_number, _) = packet_number_len
            .decode_truncated_packet_number(DecoderBuffer::new(packet_number_bytes))?;
        packet_number
    };

    Ok((
        packet_number,
        EncryptedPayload::new(header_len, packet_number_len, payload),
    ))
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::crypto;
use hex_literal::hex;

/// Types for which are able to perform initial cryptography.
///
/// This marker trait ensures only Initial-level keys
/// are used with Initial packets. Any key misuses are
/// caught by the type system.
pub trait InitialKey: crypto::Key + Sized {
    type HeaderKey: crypto::HeaderKey;

    fn new_server(connection_id: &[u8]) -> (Self, Self::HeaderKey);
    fn new_client(connection_id: &[u8]) -> (Self, Self::HeaderKey);
}

/// Types for which are able to perform initial header cryptography.
///
/// This marker trait ensures only Initial-level header keys
/// are used with Initial packets. Any key misuses are
/// caught by the type system.
pub trait InitialHeaderKey: crypto::HeaderKey {}

//= https://www.rfc-editor.org/rfc/rfc9001#section-5.2
//# initial_salt = 0x38762cf7f55934b34d179ae6a4c80cadccbb7f0a

pub const INITIAL_SALT: [u8; 20] = hex!("38762cf7f55934b34d179ae6a4c80cadccbb7f0a");

//= https://www.rfc-editor.org/rfc/rfc9001#section-5.2
//# client_initial_secret = HKDF-Expand-Label(initial_secret,
//#                                           "client in", "",
//#                                           Hash.length)

pub const INITIAL_CLIENT_LABEL: [u8; 9] = *b"client in";

//= https://www.rfc-editor.org/rfc/rfc9001#section-5.2
//# server_initial_secret = HKDF-Expand-Label(initial_secret,
//#                                           "server in", "",
//#                                           Hash.length)

pub const INITIAL_SERVER_LABEL: [u8; 9] = *b"server in";

//= https://www.rfc-editor.org/rfc/rfc9001#appendix-A
//# These packets use an 8-byte client-chosen Destination Connection ID
//# of 0x8394c8f03e515708.

pub const EXAMPLE_DCID: [u8; 8] = hex!("8394c8f03e515708");

//= https://www.rfc-editor.org/rfc/rfc9001#appendix-A.1
//# client_initial_secret
//#     = HKDF-Expand-Label(initial_secret, "client in", "", 32)
//#     = c00cf151ca5be075ed0ebfb5c80323c4
//#       2d6b7db67881289af4008f1f6c357aea

pub const EXAMPLE_CLIENT_INITIAL_SECRET: [u8; 32] = hex!(
    "
    c00cf151ca5be075ed0ebfb5c80323c4
    2d6b7db67881289af4008f1f6c357aea
    "
);

//= https://www.rfc-editor.org/rfc/rfc9001#appendix-A.1
//# server_initial_secret
//#     = HKDF-Expand-Label(initial_secret, "server in", "", 32)
//#     = 3c199828fd139efd216c155ad844cc81
//#       fb82fa8d7446fa7d78be803acdda951b

pub const EXAMPLE_SERVER_INITIAL_SECRET: [u8; 32] = hex!(
    "
    3c199828fd139efd216c155ad844cc81
    fb82fa8d7446fa7d78be803acdda951b
    "
);

//= https://www.rfc-editor.org/rfc/rfc9001#appendix-A.2
//# The client sends an Initial packet.  The unprotected payload of this
//# packet contains the following CRYPTO frame, plus enough PADDING
//# frames to make an 1162-byte payload:
//#
//# 060040f1010000ed0303ebf8fa56f129 39b9584a3896472ec40bb863cfd3e868
//# 04fe3a47f06a2b69484c000004130113 02010000c000000010000e00000b6578
//# 616d706c652e636f6dff01000100000a 00080006001d00170018001000070005
//# 04616c706e0005000501000000000033 00260024001d00209370b2c9caa47fba
//# baf4559fedba753de171fa71f50f1ce1 5d43e994ec74d748002b000302030400
//# 0d0010000e0403050306030203080408 050806002d00020101001c0002400100
//# 3900320408ffffffffffffffff050480 00ffff07048000ffff08011001048000
//# 75300901100f088394c8f03e51570806 048000ffff

/// Example payload from <https://www.rfc-editor.org/rfc/rfc9001#appendix-A.2>
pub const EXAMPLE_CLIENT_INITIAL_PAYLOAD: [u8; 245] = hex!(
    "
   060040f1010000ed0303ebf8fa56f129 39b9584a3896472ec40bb863cfd3e868
   04fe3a47f06a2b69484c000004130113 02010000c000000010000e00000b6578
   616d706c652e636f6dff01000100000a 00080006001d00170018001000070005
   04616c706e0005000501000000000033 00260024001d00209370b2c9caa47fba
   baf4559fedba753de171fa71f50f1ce1 5d43e994ec74d748002b000302030400
   0d0010000e0403050306030203080408 050806002d00020101001c0002400100
   3900320408ffffffffffffffff050480 00ffff07048000ffff08011001048000
   75300901100f088394c8f03e51570806 048000ffff
    "
);

//= https://www.rfc-editor.org/rfc/rfc9001#appendix-A.2
//# The unprotected header indicates a length of 1182 bytes: the 4-byte
//# packet number, 1162 bytes of frames, and the 16-byte authentication
//# tag.  The header includes the connection ID and a packet number of 2:
//#
//# c300000001088394c8f03e5157080000449e00000002

pub const EXAMPLE_CLIENT_INITIAL_HEADER: [u8; 22] =
    hex!("c300000001088394c8f03e5157080000449e00000002");

//= https://www.rfc-editor.org/rfc/rfc9001#appendix-A.2
//# Protecting the payload produces output that is sampled for header
//# protection.  Because the header uses a 4-byte packet number encoding,
//# the first 16 bytes of the protected payload is sampled and then
//# applied to the header as follows:
//#
//# sample = d1b1c98dd7689fb8ec11d242b123dc9b
//#
//# mask = AES-ECB(hp, sample)[0..4]
//#      = 437b9aec36
//#
//# header[0] ^= mask[0] & 0x0f
//#      = c0
//# header[18..21] ^= mask[1..4]
//#      = 7b9aec34
//# header = c000000001088394c8f03e5157080000449e7b9aec34

#[test]
fn client_initial_protection_test() {
    let mask = hex!("437b9aec36");
    let unprotected_header = EXAMPLE_CLIENT_INITIAL_HEADER;
    let protected_header = hex!("c000000001088394c8f03e5157080000449e7b9aec34");
    let packet_tag = 0b11; // results in 4 byte packet number

    header_protection_test_helper(mask, &unprotected_header, &protected_header, packet_tag);
}

//= https://www.rfc-editor.org/rfc/rfc9001#appendix-A.2
//# The resulting protected packet is:
//#
//# c000000001088394c8f03e5157080000 449e7b9aec34d1b1c98dd7689fb8ec11
//# d242b123dc9bd8bab936b47d92ec356c 0bab7df5976d27cd449f63300099f399
//# 1c260ec4c60d17b31f8429157bb35a12 82a643a8d2262cad67500cadb8e7378c
//# 8eb7539ec4d4905fed1bee1fc8aafba1 7c750e2c7ace01e6005f80fcb7df6212
//# 30c83711b39343fa028cea7f7fb5ff89 eac2308249a02252155e2347b63d58c5
//# 457afd84d05dfffdb20392844ae81215 4682e9cf012f9021a6f0be17ddd0c208
//# 4dce25ff9b06cde535d0f920a2db1bf3 62c23e596d11a4f5a6cf3948838a3aec
//# 4e15daf8500a6ef69ec4e3feb6b1d98e 610ac8b7ec3faf6ad760b7bad1db4ba3
//# 485e8a94dc250ae3fdb41ed15fb6a8e5 eba0fc3dd60bc8e30c5c4287e53805db
//# 059ae0648db2f64264ed5e39be2e20d8 2df566da8dd5998ccabdae053060ae6c
//# 7b4378e846d29f37ed7b4ea9ec5d82e7 961b7f25a9323851f681d582363aa5f8
//# 9937f5a67258bf63ad6f1a0b1d96dbd4 faddfcefc5266ba6611722395c906556
//# be52afe3f565636ad1b17d508b73d874 3eeb524be22b3dcbc2c7468d54119c74
//# 68449a13d8e3b95811a198f3491de3e7 fe942b330407abf82a4ed7c1b311663a
//# c69890f4157015853d91e923037c227a 33cdd5ec281ca3f79c44546b9d90ca00
//# f064c99e3dd97911d39fe9c5d0b23a22 9a234cb36186c4819e8b9c5927726632
//# 291d6a418211cc2962e20fe47feb3edf 330f2c603a9d48c0fcb5699dbfe58964
//# 25c5bac4aee82e57a85aaf4e2513e4f0 5796b07ba2ee47d80506f8d2c25e50fd
//# 14de71e6c418559302f939b0e1abd576 f279c4b2e0feb85c1f28ff18f58891ff
//# ef132eef2fa09346aee33c28eb130ff2 8f5b766953334113211996d20011a198
//# e3fc433f9f2541010ae17c1bf202580f 6047472fb36857fe843b19f5984009dd
//# c324044e847a4f4a0ab34f719595de37 252d6235365e9b84392b061085349d73
//# 203a4a13e96f5432ec0fd4a1ee65accd d5e3904df54c1da510b0ff20dcc0c77f
//# cb2c0e0eb605cb0504db87632cf3d8b4 dae6e705769d1de354270123cb11450e
//# fc60ac47683d7b8d0f811365565fd98c 4c8eb936bcab8d069fc33bd801b03ade
//# a2e1fbc5aa463d08ca19896d2bf59a07 1b851e6c239052172f296bfb5e724047
//# 90a2181014f3b94a4e97d117b4381303 68cc39dbb2d198065ae3986547926cd2
//# 162f40a29f0c3c8745c0f50fba3852e5 66d44575c29d39a03f0cda721984b6f4
//# 40591f355e12d439ff150aab7613499d bd49adabc8676eef023b15b65bfc5ca0
//# 6948109f23f350db82123535eb8a7433 bdabcb909271a6ecbcb58b936a88cd4e
//# 8f2e6ff5800175f113253d8fa9ca8885 c2f552e657dc603f252e1a8e308f76f0
//# be79e2fb8f5d5fbbe2e30ecadd220723 c8c0aea8078cdfcb3868263ff8f09400
//# 54da48781893a7e49ad5aff4af300cd8 04a6b6279ab3ff3afb64491c85194aab
//# 760d58a606654f9f4400e8b38591356f bf6425aca26dc85244259ff2b19c41b9
//# f96f3ca9ec1dde434da7d2d392b905dd f3d1f9af93d1af5950bd493f5aa731b4
//# 056df31bd267b6b90a079831aaf579be 0a39013137aac6d404f518cfd4684064
//# 7e78bfe706ca4cf5e9c5453e9f7cfd2b 8b4c8d169a44e55c88d4a9a7f9474241
//# e221af44860018ab0856972e194cd934

/// <https://www.rfc-editor.org/rfc/rfc9001#appendix-A.2>
pub const EXAMPLE_CLIENT_INITIAL_PROTECTED_PACKET: [u8; 1200] = hex!(
    "
   c000000001088394c8f03e5157080000 449e7b9aec34d1b1c98dd7689fb8ec11
   d242b123dc9bd8bab936b47d92ec356c 0bab7df5976d27cd449f63300099f399
   1c260ec4c60d17b31f8429157bb35a12 82a643a8d2262cad67500cadb8e7378c
   8eb7539ec4d4905fed1bee1fc8aafba1 7c750e2c7ace01e6005f80fcb7df6212
   30c83711b39343fa028cea7f7fb5ff89 eac2308249a02252155e2347b63d58c5
   457afd84d05dfffdb20392844ae81215 4682e9cf012f9021a6f0be17ddd0c208
   4dce25ff9b06cde535d0f920a2db1bf3 62c23e596d11a4f5a6cf3948838a3aec
   4e15daf8500a6ef69ec4e3feb6b1d98e 610ac8b7ec3faf6ad760b7bad1db4ba3
   485e8a94dc250ae3fdb41ed15fb6a8e5 eba0fc3dd60bc8e30c5c4287e53805db
   059ae0648db2f64264ed5e39be2e20d8 2df566da8dd5998ccabdae053060ae6c
   7b4378e846d29f37ed7b4ea9ec5d82e7 961b7f25a9323851f681d582363aa5f8
   9937f5a67258bf63ad6f1a0b1d96dbd4 faddfcefc5266ba6611722395c906556
   be52afe3f565636ad1b17d508b73d874 3eeb524be22b3dcbc2c7468d54119c74
   68449a13d8e3b95811a198f3491de3e7 fe942b330407abf82a4ed7c1b311663a
   c69890f4157015853d91e923037c227a 33cdd5ec281ca3f79c44546b9d90ca00
   f064c99e3dd97911d39fe9c5d0b23a22 9a234cb36186c4819e8b9c5927726632
   291d6a418211cc2962e20fe47feb3edf 330f2c603a9d48c0fcb5699dbfe58964
   25c5bac4aee82e57a85aaf4e2513e4f0 5796b07ba2ee47d80506f8d2c25e50fd
   14de71e6c418559302f939b0e1abd576 f279c4b2e0feb85c1f28ff18f58891ff
   ef132eef2fa09346aee33c28eb130ff2 8f5b766953334113211996d20011a198
   e3fc433f9f2541010ae17c1bf202580f 6047472fb36857fe843b19f5984009dd
   c324044e847a4f4a0ab34f719595de37 252d6235365e9b84392b061085349d73
   203a4a13e96f5432ec0fd4a1ee65accd d5e3904df54c1da510b0ff20dcc0c77f
   cb2c0e0eb605cb0504db87632cf3d8b4 dae6e705769d1de354270123cb11450e
   fc60ac47683d7b8d0f811365565fd98c 4c8eb936bcab8d069fc33bd801b03ade
   a2e1fbc5aa463d08ca19896d2bf59a07 1b851e6c239052172f296bfb5e724047
   90a2181014f3b94a4e97d117b4381303 68cc39dbb2d198065ae3986547926cd2
   162f40a29f0c3c8745c0f50fba3852e5 66d44575c29d39a03f0cda721984b6f4
   40591f355e12d439ff150aab7613499d bd49adabc8676eef023b15b65bfc5ca0
   6948109f23f350db82123535eb8a7433 bdabcb909271a6ecbcb58b936a88cd4e
   8f2e6ff5800175f113253d8fa9ca8885 c2f552e657dc603f252e1a8e308f76f0
   be79e2fb8f5d5fbbe2e30ecadd220723 c8c0aea8078cdfcb3868263ff8f09400
   54da48781893a7e49ad5aff4af300cd8 04a6b6279ab3ff3afb64491c85194aab
   760d58a606654f9f4400e8b38591356f bf6425aca26dc85244259ff2b19c41b9
   f96f3ca9ec1dde434da7d2d392b905dd f3d1f9af93d1af5950bd493f5aa731b4
   056df31bd267b6b90a079831aaf579be 0a39013137aac6d404f518cfd4684064
   7e78bfe706ca4cf5e9c5453e9f7cfd2b 8b4c8d169a44e55c88d4a9a7f9474241
   e221af44860018ab0856972e194cd934
    "
);

//= https://www.rfc-editor.org/rfc/rfc9001#appendix-A.3
//# The server sends the following payload in response, including an ACK
//# frame, a CRYPTO frame, and no PADDING frames:
//#
//# 02000000000600405a020000560303ee fce7f7b37ba1d1632e96677825ddf739
//# 88cfc79825df566dc5430b9a045a1200 130100002e00330024001d00209d3c94
//# 0d89690b84d08a60993c144eca684d10 81287c834d5311bcf32bb9da1a002b00
//# 020304
//#

/// Example payload from <https://www.rfc-editor.org/rfc/rfc9001#appendix-A.3>
pub const EXAMPLE_SERVER_INITIAL_PAYLOAD: [u8; 99] = hex!(
    "
   02000000000600405a020000560303ee fce7f7b37ba1d1632e96677825ddf739
   88cfc79825df566dc5430b9a045a1200 130100002e00330024001d00209d3c94
   0d89690b84d08a60993c144eca684d10 81287c834d5311bcf32bb9da1a002b00
   020304
    "
);

//= https://www.rfc-editor.org/rfc/rfc9001#appendix-A.3
//# The header from the server includes a new connection ID and a 2-byte
//# packet number encoding for a packet number of 1:
//#
//# c1000000010008f067a5502a4262b50040750001

pub const EXAMPLE_SERVER_INITIAL_HEADER: [u8; 20] =
    hex!("c1000000010008f067a5502a4262b50040750001");

//= https://www.rfc-editor.org/rfc/rfc9001#appendix-A.3
//# As a result, after protection, the header protection sample is taken
//# starting from the third protected byte:
//#
//# sample = 2cd0991cd25b0aac406a5816b6394100
//# mask   = 2ec0d8356a
//# header = cf000000010008f067a5502a4262b5004075c0d9

#[test]
fn server_initial_protection_test() {
    let mask = hex!("2ec0d8356a");

    let unprotected_header = EXAMPLE_SERVER_INITIAL_HEADER;
    let protected_header = hex!("cf000000010008f067a5502a4262b5004075c0d9");
    let packet_tag = 0b01; // results in a 2 byte packet number

    header_protection_test_helper(mask, &unprotected_header, &protected_header, packet_tag);
}

//= https://www.rfc-editor.org/rfc/rfc9001#appendix-A.3
//# The final protected packet is then:
//#
//# cf000000010008f067a5502a4262b500 4075c0d95a482cd0991cd25b0aac406a
//# 5816b6394100f37a1c69797554780bb3 8cc5a99f5ede4cf73c3ec2493a1839b3
//# dbcba3f6ea46c5b7684df3548e7ddeb9 c3bf9c73cc3f3bded74b562bfb19fb84
//# 022f8ef4cdd93795d77d06edbb7aaf2f 58891850abbdca3d20398c276456cbc4
//# 2158407dd074ee

/// Example protected packet from
/// <https://www.rfc-editor.org/rfc/rfc9001#appendix-A.3>
pub const EXAMPLE_SERVER_INITIAL_PROTECTED_PACKET: [u8; 135] = hex!(
    "
   cf000000010008f067a5502a4262b500 4075c0d95a482cd0991cd25b0aac406a
   5816b6394100f37a1c69797554780bb3 8cc5a99f5ede4cf73c3ec2493a1839b3
   dbcba3f6ea46c5b7684df3548e7ddeb9 c3bf9c73cc3f3bded74b562bfb19fb84
   022f8ef4cdd93795d77d06edbb7aaf2f 58891850abbdca3d20398c276456cbc4
   2158407dd074ee
    "
);

#[cfg(test)]
fn header_protection_test_helper(
    mask: crate::crypto::HeaderProtectionMask,
    unprotected_header: &[u8],
    protected_header: &[u8],
    packet_tag: u8,
) {
    use crate::{
        crypto::{
            apply_header_protection, remove_header_protection, EncryptedPayload, ProtectedPayload,
        },
        packet::number::PacketNumberSpace,
    };
    let space = PacketNumberSpace::Initial;

    let packet_number_len = space.new_packet_number_len(packet_tag);
    let header_len = protected_header.len() - packet_number_len.bytesize();

    let mut subject = protected_header.to_vec();

    remove_header_protection(space, mask, ProtectedPayload::new(header_len, &mut subject)).unwrap();

    assert_eq!(
        unprotected_header,
        &subject[..],
        "packet protection removal failed"
    );

    apply_header_protection(
        mask,
        EncryptedPayload::new(header_len, packet_number_len, &mut subject),
    );

    assert_eq!(
        protected_header,
        &subject[..],
        "packet protection application failed"
    );
}