One-KVM/libs/dtls/src/fragment_buffer/fragment_buffer_test.rs

use super::*;

#[test]
fn test_fragment_buffer() -> Result<()> {
    let tests = vec![
        (
            "Single Fragment",
            vec![vec![
                0x16, 0xfe, 0xff, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0F, 0x03,
                0x00, 0x00, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03, 0xfe, 0xff, 0x00,
            ]],
            vec![vec![
                0x03, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03, 0xfe, 0xff,
                0x00,
            ]],
            0,
        ),
        (
            "Single Fragment Epoch 3",
            vec![vec![
                0x16, 0xfe, 0xff, 0x00, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0F, 0x03,
                0x00, 0x00, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03, 0xfe, 0xff, 0x00,
            ]],
            vec![vec![
                0x03, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03, 0xfe, 0xff,
                0x00,
            ]],
            3,
        ),
        (
            "Multiple Fragments",
            vec![
                vec![
                    0x16, 0xfe, 0xfd, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x81,
                    0x0b, 0x00, 0x00, 0x0F, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x05, 0x00,
                    0x01, 0x02, 0x03, 0x04,
                ],
                vec![
                    0x16, 0xfe, 0xfd, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x81,
                    0x0b, 0x00, 0x00, 0x0F, 0x00, 0x00, 0x00, 0x00, 0x05, 0x00, 0x00, 0x05, 0x05,
                    0x06, 0x07, 0x08, 0x09,
                ],
                vec![
                    0x16, 0xfe, 0xfd, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x81,
                    0x0b, 0x00, 0x00, 0x0F, 0x00, 0x00, 0x00, 0x00, 0x0A, 0x00, 0x00, 0x05, 0x0A,
                    0x0B, 0x0C, 0x0D, 0x0E,
                ],
            ],
            vec![vec![
                0x0b, 0x00, 0x00, 0x0f, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0f, 0x00, 0x01,
                0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e,
            ]],
            0,
        ),
        (
            "Multiple Unordered Fragments",
            vec![
                vec![
                    0x16, 0xfe, 0xfd, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x81,
                    0x0b, 0x00, 0x00, 0x0F, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x05, 0x00,
                    0x01, 0x02, 0x03, 0x04,
                ],
                vec![
                    0x16, 0xfe, 0xfd, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x81,
                    0x0b, 0x00, 0x00, 0x0F, 0x00, 0x00, 0x00, 0x00, 0x0A, 0x00, 0x00, 0x05, 0x0A,
                    0x0B, 0x0C, 0x0D, 0x0E,
                ],
                vec![
                    0x16, 0xfe, 0xfd, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0x00, 0x81,
                    0x0b, 0x00, 0x00, 0x0F, 0x00, 0x00, 0x00, 0x00, 0x05, 0x00, 0x00, 0x05, 0x05,
                    0x06, 0x07, 0x08, 0x09,
                ],
            ],
            vec![vec![
                0x0b, 0x00, 0x00, 0x0f, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0f, 0x00, 0x01,
                0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e,
            ]],
            0,
        ),
        (
            "Multiple Handshakes in Single Fragment",
            vec![vec![
                0x16, 0xfe, 0xfd, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00,
                0x30, /* record header */
                0x03, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x04, 0xfe, 0xff,
                0x01, 0x01, /*handshake msg 1*/
                0x03, 0x00, 0x00, 0x04, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x04, 0xfe, 0xff,
                0x01, 0x01, /*handshake msg 2*/
                0x03, 0x00, 0x00, 0x04, 0x00, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x04, 0xfe, 0xff,
                0x01, 0x01, /*handshake msg 3*/
            ]],
            vec![
                vec![
                    0x03, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x04, 0xfe,
                    0xff, 0x01, 0x01,
                ],
                vec![
                    0x03, 0x00, 0x00, 0x04, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x04, 0xfe,
                    0xff, 0x01, 0x01,
                ],
                vec![
                    0x03, 0x00, 0x00, 0x04, 0x00, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x04, 0xfe,
                    0xff, 0x01, 0x01,
                ],
            ],
            0,
        ),
        // Ensure zero length fragments don't cause an infinite recursive loop which in turn causes
        // a stack overflow.
        (
            "Zero length fragment",
            vec![vec![
                0x16, 0xfe, 0xfd, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0d, 0x00,
                0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            ]],
            vec![vec![
                0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00,
            ]],
            0,
        ),
    ];

    for (name, inputs, expects, expected_epoch) in tests {
        let mut fragment_buffer = FragmentBuffer::new();
        for frag in inputs {
            let status = fragment_buffer.push(&frag)?;
            assert!(
                status,
                "fragment_buffer didn't accept fragments for '{name}'"
            );
        }

        for expected in expects {
            let (out, epoch) = fragment_buffer.pop()?;
            assert_eq!(
                out, expected,
                "fragment_buffer '{name}' push/pop: got {out:?}, want {expected:?}"
            );

            assert_eq!(
                epoch, expected_epoch,
                "fragment_buffer returned wrong epoch: got {epoch}, want {expected_epoch}"
            );
        }

        let result = fragment_buffer.pop();
        assert!(
            result.is_err(),
            "fragment_buffer popped single buffer multiple times for '{name}'"
        );
    }

    Ok(())
}

#[test]
fn test_fragment_buffer_overflow() -> Result<()> {
    let mut fragment_buffer = FragmentBuffer::new();

    fragment_buffer.push(&[
        0x16, 0xfe, 0xff, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0F, 0x03, 0x00,
        0x00, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03, 0xfe, 0xff, 0x00,
    ])?;

    let big_buffer = vec![0; 2_000_000];
    let result = fragment_buffer.push(&big_buffer);

    assert!(
        result.is_err(),
        "Pushing a buffer of size 2MB should have caused FragmentBuffer::push to return an error"
    );

    Ok(())
}