#![no_main]

use commonware_codec::FixedSize;
use commonware_cryptography::{ChaCha20Poly1305, Cipher, Signer, ed25519::PrivateKey};
use commonware_runtime::{
    Handle, Runner as _, Sink as _, Spawner, Stream as _, Supervisor as _, deterministic, mocks,
};
use commonware_stream::{
    SakeCups, Upgrader as StreamUpgrader,
    cups::{self, Cups},
    sake::{self, Sake},
    utils::{
        Timeout,
        codec::{Error, recv_frame, send_frame},
    },
};
use futures::future::{Either, select};
use libfuzzer_sys::fuzz_target;
use std::time::Duration;

const NAMESPACE: &[u8] = b"fuzz_transport";
const MAX_MESSAGE_SIZE: u32 = 2048;
const MAX_CIPHERTEXT_SIZE: u32 =
    MAX_MESSAGE_SIZE + <<ChaCha20Poly1305 as Cipher>::Tag as FixedSize>::SIZE as u32;

#[derive(Debug, arbitrary::Arbitrary)]
enum Direction {
    D2L,
    L2D,
}

#[derive(Debug)]
enum Message {
    Authenticated(Direction, Vec<u8>),
    Unauthenticated(Direction, Vec<u8>),
    Tampered(Direction, Vec<u8>, usize, u8),
}

impl<'a> arbitrary::Arbitrary<'a> for Message {
    fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
        let direction = Direction::arbitrary(u)?;
        let kind = u.int_in_range(0..=2)?;

        // Forged records may be as long as a ciphertext. The other kinds carry plaintext.
        let max_len = if kind == 1 {
            MAX_CIPHERTEXT_SIZE as usize
        } else {
            MAX_MESSAGE_SIZE as usize
        };
        let len = u.int_in_range(0..=max_len.min(u.len()))?;
        let msg = u.bytes(len)?.to_vec();
        let out = match kind {
            0 => Self::Authenticated(direction, msg),
            1 => Self::Unauthenticated(direction, msg),

            // A zero mask would leave the record unchanged.
            _ => Self::Tampered(direction, msg, u.arbitrary()?, u.arbitrary::<u8>()?.max(1)),
        };
        Ok(out)
    }
}

/// Sending half of an [Upgrader] connection over mock channels.
type Sender =
    <SakeCups<PrivateKey, ChaCha20Poly1305> as StreamUpgrader>::Sender<mocks::Stream, mocks::Sink>;

/// Receiving half of an [Upgrader] connection over mock channels.
type Receiver = <SakeCups<PrivateKey, ChaCha20Poly1305> as StreamUpgrader>::Receiver<
    mocks::Stream,
    mocks::Sink,
>;

#[derive(Debug)]
pub struct FuzzInput {
    sake: sake::Version,
    cups: cups::Version,
    setup_corruption: Vec<u8>,
    messages: Vec<Message>,
}

impl FuzzInput {
    fn has_setup_corruption(&self) -> bool {
        self.setup_corruption.iter().any(|&x| x != 0)
    }
}

impl<'a> arbitrary::Arbitrary<'a> for FuzzInput {
    fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
        let sake = u.arbitrary()?;
        let cups = u.arbitrary()?;
        let setup_corruption = if bool::arbitrary(u)? {
            Vec::arbitrary(u)?
        } else {
            Vec::new()
        };
        let messages = u.arbitrary_iter()?.collect::<Result<Vec<Message>, _>>()?;
        Ok(Self {
            sake,
            cups,
            setup_corruption,
            messages,
        })
    }
}

fn fuzz(input: FuzzInput) {
    let executor = deterministic::Runner::default();
    executor.start(|context| async move {
        let has_setup_corruption = input.has_setup_corruption();
        let FuzzInput {
            sake,
            cups,
            setup_corruption,
            messages,
        } = input;
        let layout = Cups::<(), ChaCha20Poly1305>::new((), cups);
        let dialer_signer = PrivateKey::from_seed(42);
        let listener_signer = PrivateKey::from_seed(24);

        let (dialer_sink, mut adversary_d_stream) = mocks::Channel::init();
        let (mut adversary_d_sink, listener_stream) = mocks::Channel::init();
        let (listener_sink, mut adversary_l_stream) = mocks::Channel::init();
        let (mut adversary_l_sink, dialer_stream) = mocks::Channel::init();

        let dialer_handshake = Cups::<_, ChaCha20Poly1305>::new(
            Sake {
                signer: dialer_signer.clone(),
                synchrony_bound: Duration::from_secs(1),
                max_handshake_age: Duration::from_secs(1),
                version: sake,
            },
            cups,
        );
        let dialer_handshake = Timeout::new(dialer_handshake, Duration::from_secs(1));

        let listener_handshake = Cups::<_, ChaCha20Poly1305>::new(
            Sake {
                signer: listener_signer.clone(),
                synchrony_bound: Duration::from_secs(1),
                max_handshake_age: Duration::from_secs(1),
                version: sake,
            },
            cups,
        );
        let listener_handshake = Timeout::new(listener_handshake, Duration::from_secs(1));

        let dialer_handle = context.child("dialer").spawn(move |context| async move {
            dialer_handshake
                .dial(
                    context,
                    NAMESPACE,
                    MAX_MESSAGE_SIZE,
                    listener_signer.public_key(),
                    dialer_stream,
                    dialer_sink,
                )
                .await
        });
        let listener_handle = context.child("listener").spawn(move |context| async move {
            listener_handshake
                .listen(
                    context,
                    NAMESPACE,
                    MAX_MESSAGE_SIZE,
                    |_| async { true },
                    listener_stream,
                    listener_sink,
                )
                .await
        });
        let adversary_handle: Handle<Result<_, Error>> =
            context
                .child("adversary")
                .spawn(move |_context| async move {
                    let mut corruption_i = 0;

                    let announce = recv_frame(&mut adversary_d_stream, MAX_MESSAGE_SIZE).await?;
                    send_frame(&mut adversary_d_sink, announce, MAX_MESSAGE_SIZE).await?;

                    let mut m1: Vec<u8> = recv_frame(&mut adversary_d_stream, MAX_MESSAGE_SIZE)
                        .await?
                        .coalesce()
                        .into();
                    for byte in m1.iter_mut() {
                        if corruption_i < setup_corruption.len() {
                            *byte ^= setup_corruption[corruption_i];
                            corruption_i += 1;
                        }
                    }
                    send_frame(&mut adversary_d_sink, m1, MAX_MESSAGE_SIZE).await?;

                    let mut m2: Vec<u8> = recv_frame(&mut adversary_l_stream, MAX_MESSAGE_SIZE)
                        .await?
                        .coalesce()
                        .into();
                    for byte in m2.iter_mut() {
                        if corruption_i < setup_corruption.len() {
                            *byte ^= setup_corruption[corruption_i];
                            corruption_i += 1;
                        }
                    }
                    send_frame(&mut adversary_l_sink, m2, MAX_MESSAGE_SIZE).await?;

                    let mut m3: Vec<u8> = recv_frame(&mut adversary_d_stream, MAX_MESSAGE_SIZE)
                        .await?
                        .coalesce()
                        .into();
                    for byte in m3.iter_mut() {
                        if corruption_i < setup_corruption.len() {
                            *byte ^= setup_corruption[corruption_i];
                            corruption_i += 1;
                        }
                    }
                    let sent_corrupted_data =
                        setup_corruption.iter().take(corruption_i).any(|x| *x != 0);
                    send_frame(&mut adversary_d_sink, m3, MAX_MESSAGE_SIZE).await?;
                    Ok((
                        sent_corrupted_data,
                        adversary_d_stream,
                        adversary_d_sink,
                        adversary_l_stream,
                        adversary_l_sink,
                    ))
                });

        // We need to do a selection to correctly assert the errors, avoiding deadlock.
        //
        // A deadlock might happen if one side asserts an error, and then we're foolishly waiting
        // for it to send a message it never will.
        let (d_res, l_res) = match select(dialer_handle, listener_handle).await {
            Either::Left((d_res, l_handle)) => {
                match d_res.inspect_err(|e| println!("A: {e:?}")).unwrap().ok() {
                    Some(d_res) => (Some(d_res), l_handle.await.unwrap().ok()),
                    None => (None, None),
                }
            }
            Either::Right((l_res, d_handle)) => {
                match l_res.inspect_err(|e| println!("B: {e:?}")).unwrap().ok() {
                    Some(l_res) => (d_handle.await.unwrap().ok(), Some(l_res)),
                    None => (None, None),
                }
            }
        };
        if d_res.is_none() || l_res.is_none() {
            assert!(has_setup_corruption, "expected there to be data corruption");
            return;
        }
        let (mut d_sender, mut d_receiver) = d_res.unwrap();
        let (_, mut l_sender, mut l_receiver) = l_res.unwrap();
        let (
            sent_corrupted_data,
            mut adversary_d_stream,
            mut adversary_d_sink,
            mut adversary_l_stream,
            mut adversary_l_sink,
        ) = adversary_handle.await.unwrap().unwrap();
        // Importantly, make sure that if we've gotten to this point, no data corruption
        // has happened!
        assert!(!sent_corrupted_data);
        let mut d2l_corrupted = false;
        let mut l2d_corrupted = false;
        for msg in messages {
            match msg {
                Message::Authenticated(direction, data) => {
                    if matches!(&direction, Direction::D2L) && d2l_corrupted {
                        continue;
                    }
                    if matches!(&direction, Direction::L2D) && l2d_corrupted {
                        continue;
                    }
                    let (sender, a_in, a_out, receiver): (
                        &mut Sender,
                        &mut mocks::Stream,
                        &mut mocks::Sink,
                        &mut Receiver,
                    ) = match direction {
                        Direction::D2L => (
                            &mut d_sender,
                            &mut adversary_d_stream,
                            &mut adversary_d_sink,
                            &mut l_receiver,
                        ),
                        Direction::L2D => (
                            &mut l_sender,
                            &mut adversary_l_stream,
                            &mut adversary_l_sink,
                            &mut d_receiver,
                        ),
                    };

                    // Send a legitimate plaintext message through the connection.
                    sender.send(data.clone()).await.unwrap();

                    // Intercept the resulting record from the wire.
                    let record = a_in.recv(layout.record_len(data.len())).await.unwrap();

                    // Forward the exact record unchanged.
                    a_out.send(record).await.unwrap();

                    // Receiver should decrypt and deliver the original plaintext.
                    let data2 = receiver.recv().await.unwrap();
                    assert_eq!(data2.coalesce(), data.as_slice(), "expected data to match");
                }
                Message::Unauthenticated(direction, data) => {
                    if matches!(&direction, Direction::D2L) && d2l_corrupted {
                        continue;
                    }
                    if matches!(&direction, Direction::L2D) && l2d_corrupted {
                        continue;
                    }
                    let (sender, a_in, a_out, receiver): (
                        &mut Sender,
                        &mut mocks::Stream,
                        &mut mocks::Sink,
                        &mut Receiver,
                    ) = match direction {
                        Direction::D2L => (
                            &mut d_sender,
                            &mut adversary_d_stream,
                            &mut adversary_d_sink,
                            &mut l_receiver,
                        ),
                        Direction::L2D => (
                            &mut l_sender,
                            &mut adversary_l_stream,
                            &mut adversary_l_sink,
                            &mut d_receiver,
                        ),
                    };

                    // Send one legitimate record so the sender's cipher advances.
                    sender.send(vec![0u8]).await.unwrap();

                    // Adversary intercepts and drops that record.
                    let _ = a_in.recv(layout.record_len(1)).await.unwrap();

                    // Adversary injects forged unauthenticated bytes instead. A forged version 1
                    // header is padded to full size so the receiver has a header to reject.
                    match cups {
                        cups::Version::V0 => {
                            send_frame(a_out, data, MAX_CIPHERTEXT_SIZE).await.unwrap()
                        }
                        cups::Version::V1 => {
                            let mut forged = data;
                            forged.resize(forged.len().max(layout.header_len(0)), 0);
                            a_out.send(forged).await.unwrap();
                        }
                    }

                    // Receiver must reject the forged record.
                    let res = receiver.recv().await;
                    assert!(res.is_err());

                    // After a failed open, this direction's receiver refuses every later record.
                    if matches!(&direction, Direction::D2L) {
                        d2l_corrupted = true;
                    } else {
                        l2d_corrupted = true;
                    }
                }
                Message::Tampered(direction, data, index, mask) => {
                    if matches!(&direction, Direction::D2L) && d2l_corrupted {
                        continue;
                    }
                    if matches!(&direction, Direction::L2D) && l2d_corrupted {
                        continue;
                    }
                    let (sender, a_in, a_out, receiver): (
                        &mut Sender,
                        &mut mocks::Stream,
                        &mut mocks::Sink,
                        &mut Receiver,
                    ) = match direction {
                        Direction::D2L => (
                            &mut d_sender,
                            &mut adversary_d_stream,
                            &mut adversary_d_sink,
                            &mut l_receiver,
                        ),
                        Direction::L2D => (
                            &mut l_sender,
                            &mut adversary_l_stream,
                            &mut adversary_l_sink,
                            &mut d_receiver,
                        ),
                    };

                    // Send a legitimate record and intercept it.
                    sender.send(data.clone()).await.unwrap();
                    let mut record: Vec<u8> = a_in
                        .recv(layout.record_len(data.len()))
                        .await
                        .unwrap()
                        .coalesce()
                        .into();

                    // Flip one byte. A version 0 prefix is left intact so the receiver reads the
                    // same span, and any version 1 byte may change because its header is
                    // authenticated before the payload is requested.
                    let start = match cups {
                        cups::Version::V0 => layout.header_len(data.len()),
                        cups::Version::V1 => 0,
                    };
                    let target = start + index % (record.len() - start);
                    record[target] ^= mask;
                    a_out.send(record).await.unwrap();

                    // Receiver must reject the tampered record.
                    let res = receiver.recv().await;
                    assert!(res.is_err());

                    // After a failed open, this direction's receiver refuses every later record.
                    if matches!(&direction, Direction::D2L) {
                        d2l_corrupted = true;
                    } else {
                        l2d_corrupted = true;
                    }
                }
            }
        }
    });
}

fuzz_target!(|input: FuzzInput| {
    fuzz(input);
});