//! Property tests for the wire codec.
//!
//! Two things are checked for every kind of message:
//!
//! - **Round trip.** A message that can be written can be read back, comes
//!   back equal.
//! - **Hostile input.** Whatever bytes arrive, decoding returns a message or
//!   an error: it never panics. What it does accept can be shown, written
//!   again, and read back as the same message.
//!
//! The second is what a peer can do to us, so it is run on three kinds of
//! input: plain noise, noise behind a valid header (so it reaches the body
//! parsers), and valid messages with a few bytes damaged (so it reaches deep
//! into them).

use std::net::{Ipv4Addr, Ipv6Addr};

use bgp_packet::errors::BgpError;
use bgp_packet::message::{
    AggregatorPathAttribute, AsPathAttribute, AsPathSegment, AtomicAggregatePathAttribute,
    BGP_HEADER_LEN, BGP_MAX_MSG_LEN, CommunitiesPathAttribute, ExtendedCommunitiesPathAttribute,
    ExtendedCommunity, LargeCommunitiesPathAttribute, LargeCommunity, LocalPrefPathAttribute,
    Message, MpReachNlriPathAttribute, MpUnreachNlriPathAttribute, MultiExitDiscPathAttribute,
    NextHopPathAttribute, NlriNextHop, NotificationMessage, OriginPathAttribute, PathAttribute,
    PathAttributeFlags, RouteRefreshMessage, UpdateMessage,
};
use bgp_packet::open::{
    AddPathEntry, AddPathSendRecv, BgpRole, BgpSecHeader, Capability, ExtendedNextHop,
    GracefulRestartEntry, LongLivedGracefulRestart, MultiLabelEntry, OpenMessage, OpenOption,
    OrfFamily, OrfSendRecv, OrfType, OutboundRouteFilterEntry,
};
use bgp_primitives::constants::{Afi, AfiSafi, IpVersion, Safi};
use bgp_primitives::ip_prefix::{IpPrefix, Nlri};
use bgp_primitives::parser::{EncodeError, ParserContext};
use bytes::BytesMut;
use proptest::collection::vec;
use proptest::prelude::*;

/// What a session settled that the codec depends on.
#[derive(Debug, Clone, Copy)]
struct Session {
    four_octet_asn: bool,
    ibgp: bool,
    /// Whether prefixes carry path identifiers (RFC 7911), for IPv4 and
    /// for IPv6 unicast.
    add_path: [bool; 2],
}

impl Session {
    fn ctx(self) -> ParserContext {
        ParserContext {
            four_octet_asn: self.four_octet_asn,
            ibgp: self.ibgp,
            add_path: [AfiSafi::IPV4_UNICAST, AfiSafi::IPV6_UNICAST]
                .into_iter()
                .zip(self.add_path)
                .filter_map(|(family, in_use)| in_use.then_some(family))
                .collect(),
        }
    }

    fn has_add_path(self, version: IpVersion, safi: Safi) -> bool {
        safi == Safi::UNICAST
            && match version {
                IpVersion::Ipv4 => self.add_path[0],
                IpVersion::Ipv6 => self.add_path[1],
            }
    }
}

fn session() -> impl Strategy<Value = Session> {
    any::<(bool, bool, [bool; 2])>().prop_map(|(four_octet_asn, ibgp, add_path)| Session {
        four_octet_asn,
        ibgp,
        add_path,
    })
}

// ---------------------------------------------------------------------------
// Generators
// ---------------------------------------------------------------------------

fn prefix(version: IpVersion) -> impl Strategy<Value = IpPrefix> {
    let max_len = match version {
        IpVersion::Ipv4 => 32u8,
        IpVersion::Ipv6 => 128,
    };
    (0..=max_len, any::<[u8; 16]>()).prop_map(move |(length, bytes)| {
        IpPrefix::new(version, bytes, length).expect("16 bytes are enough for any prefix")
    })
}

/// Prefixes as an UPDATE names them. Without ADD-PATH for the family there
/// is no identifier on the wire, and it reads back as zero.
fn nlri(session: Session, version: IpVersion, safi: Safi) -> impl Strategy<Value = Vec<Nlri>> {
    let with_path_id = session.has_add_path(version, safi);
    vec(
        (prefix(version), any::<u32>()).prop_map(move |(prefix, path_id)| Nlri {
            prefix,
            path_id: if with_path_id { path_id } else { 0 },
        }),
        0..5,
    )
}

fn afi() -> impl Strategy<Value = Afi> {
    any::<u16>().prop_map(Afi)
}

fn safi() -> impl Strategy<Value = Safi> {
    any::<u8>().prop_map(Safi)
}

/// AS numbers that fit the session: any on a four-octet one, two-octet ones
/// otherwise (wider ones are replaced by AS_TRANS on the wire, by design).
fn asn(four_octet_asn: bool) -> BoxedStrategy<u32> {
    if four_octet_asn {
        any::<u32>().boxed()
    } else {
        any::<u16>().prop_map(u32::from).boxed()
    }
}

fn as_path(asn: BoxedStrategy<u32>) -> impl Strategy<Value = AsPathAttribute> {
    // A segment has at least one AS number.
    vec((any::<bool>(), vec(asn, 1..6)), 0..4).prop_map(|segments| AsPathAttribute {
        segments: segments
            .into_iter()
            .map(|(ordered, path)| AsPathSegment { ordered, path })
            .collect(),
    })
}

fn extended_community() -> impl Strategy<Value = ExtendedCommunity> {
    any::<[u8; 8]>().prop_map(ExtendedCommunity)
}

fn mp_reach(session: Session) -> impl Strategy<Value = MpReachNlriPathAttribute> {
    let ip_safi = || prop_oneof![Just(Safi::UNICAST), Just(Safi::MULTICAST)];
    let v4 = ip_safi()
        .prop_flat_map(move |safi| {
            (
                Just(safi),
                any::<[u8; 4]>(),
                nlri(session, IpVersion::Ipv4, safi),
            )
        })
        .prop_map(|(safi, next_hop, prefixes)| MpReachNlriPathAttribute {
            family: AfiSafi::new(Afi::IPV4, safi),
            next_hop: NlriNextHop::Ipv4(Ipv4Addr::from(next_hop)),
            prefixes,
        });
    let v6_next_hop = || {
        prop_oneof![
            any::<[u8; 16]>().prop_map(|a| NlriNextHop::Ipv6(Ipv6Addr::from(a))),
            (any::<[u8; 16]>(), any::<[u8; 16]>()).prop_map(|(global, link_local)| {
                NlriNextHop::Ipv6WithLl {
                    global: Ipv6Addr::from(global),
                    link_local: Ipv6Addr::from(link_local),
                }
            }),
        ]
    };
    let v6 = ip_safi()
        .prop_flat_map(move |safi| {
            (
                Just(safi),
                v6_next_hop(),
                nlri(session, IpVersion::Ipv6, safi),
            )
        })
        .prop_map(|(safi, next_hop, prefixes)| MpReachNlriPathAttribute {
            family: AfiSafi::new(Afi::IPV6, safi),
            next_hop,
            prefixes,
        });
    prop_oneof![v4, v6]
}

fn mp_unreach(session: Session) -> impl Strategy<Value = MpUnreachNlriPathAttribute> {
    let of = move |version: IpVersion| {
        nlri(session, version, Safi::UNICAST).prop_map(move |prefixes| MpUnreachNlriPathAttribute {
            family: AfiSafi::new(version.into(), Safi::UNICAST),
            prefixes,
        })
    };
    prop_oneof![of(IpVersion::Ipv4), of(IpVersion::Ipv6)]
}

/// Type codes whose definition the decoder knows; an unknown attribute has
/// none of these.
const DEFINED_ATTRIBUTES: [u8; 16] = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 14, 15, 16, 17, 18, 32];

/// Unknown attributes, each with a type code of its own: a repeated
/// attribute is dropped by the decoder.
fn unknown_attributes() -> impl Strategy<Value = Vec<PathAttribute>> {
    proptest::collection::btree_map(
        any::<u8>().prop_filter("a type code without a definition", |code| {
            !DEFINED_ATTRIBUTES.contains(code)
        }),
        (any::<u8>(), vec(any::<u8>(), 0..300)),
        0..3,
    )
    .prop_map(|attributes| {
        attributes
            .into_iter()
            .map(|(type_code, (flags, payload))| {
                // Optional (bit 7): an unknown well-known attribute is an
                // error. Bit 4 selects the two-octet length, which a payload
                // this long needs.
                let mut flags = flags | 0x80;
                if payload.len() > 255 {
                    flags |= 0x10;
                }
                PathAttribute::UnknownPathAttribute {
                    flags: PathAttributeFlags(flags),
                    type_code,
                    payload,
                }
            })
            .collect()
    })
}

/// The attributes of a well-formed UPDATE on `session`: the mandatory ones,
/// and any of the others at most once.
fn path_attributes(session: Session) -> impl Strategy<Value = Vec<PathAttribute>> {
    let mandatory = (
        prop_oneof![
            Just(OriginPathAttribute::IGP),
            Just(OriginPathAttribute::EGP),
            Just(OriginPathAttribute::INCOMPLETE)
        ],
        as_path(asn(session.four_octet_asn)),
        any::<[u8; 4]>(),
        any::<u32>(),
    )
        .prop_map(move |(origin, as_path, next_hop, local_pref)| {
            let mut attributes = vec![
                PathAttribute::Origin(origin),
                PathAttribute::AsPath(as_path),
                PathAttribute::NextHop(NextHopPathAttribute(Ipv4Addr::from(next_hop))),
            ];
            // Mandatory between internal peers, dropped from external ones.
            if session.ibgp {
                attributes.push(PathAttribute::LocalPref(LocalPrefPathAttribute(local_pref)));
            }
            attributes
        });
    let optional = (
        proptest::option::of(
            any::<u32>().prop_map(|v| PathAttribute::MultiExitDisc(MultiExitDiscPathAttribute(v))),
        ),
        proptest::option::of(Just(PathAttribute::AtomicAggregate(
            AtomicAggregatePathAttribute {},
        ))),
        proptest::option::of((asn(session.four_octet_asn), any::<[u8; 4]>()).prop_map(
            |(asn, ip)| {
                PathAttribute::Aggregator(AggregatorPathAttribute {
                    asn,
                    ip: Ipv4Addr::from(ip),
                })
            },
        )),
        proptest::option::of(
            vec(any::<(u16, u16)>(), 1..8)
                .prop_map(|c| PathAttribute::Communities(CommunitiesPathAttribute(c))),
        ),
        proptest::option::of(mp_reach(session).prop_map(PathAttribute::MpReachNlri)),
        proptest::option::of(mp_unreach(session).prop_map(PathAttribute::MpUnreachNlri)),
        proptest::option::of(
            vec(extended_community(), 1..6).prop_map(|extended_communities| {
                PathAttribute::ExtendedCommunities(ExtendedCommunitiesPathAttribute {
                    extended_communities,
                })
            }),
        ),
        // AS4_PATH is four octets wide, and only a two-octet speaker sends
        // it.
        proptest::option::of(as_path(any::<u32>().boxed()).prop_map(PathAttribute::As4Path))
            .prop_map(move |as4_path| as4_path.filter(|_| !session.four_octet_asn)),
        proptest::option::of(vec(any::<(u32, u32, u32)>(), 1..6).prop_map(|communities| {
            PathAttribute::LargeCommunities(LargeCommunitiesPathAttribute {
                communities: communities
                    .into_iter()
                    .map(|(global_admin, data_1, data_2)| LargeCommunity {
                        global_admin,
                        data_1,
                        data_2,
                    })
                    .collect(),
            })
        })),
        unknown_attributes(),
    );
    (mandatory, optional).prop_map(|(mut attributes, optional)| {
        let (med, atomic, aggregator, communities, reach, unreach, extended, as4, large, unknown) =
            optional;
        attributes.extend(
            [
                med,
                atomic,
                aggregator,
                communities,
                reach,
                unreach,
                extended,
                as4,
                large,
            ]
            .into_iter()
            .flatten(),
        );
        attributes.extend(unknown);
        attributes
    })
}

fn update(session: Session) -> impl Strategy<Value = UpdateMessage> {
    (
        nlri(session, IpVersion::Ipv4, Safi::UNICAST),
        path_attributes(session),
        nlri(session, IpVersion::Ipv4, Safi::UNICAST),
    )
        .prop_map(
            |(withdrawn_routes, path_attributes, announced_routes)| UpdateMessage {
                withdrawn_routes,
                path_attributes,
                announced_routes,
            },
        )
}

/// Capability codes that have their own variant.
const KNOWN_CAPABILITIES: [u8; 13] = [1, 2, 3, 5, 6, 7, 8, 9, 64, 65, 69, 70, 71];

fn capability() -> impl Strategy<Value = Capability> {
    let orf = (
        any::<u8>().prop_map(OrfType),
        prop_oneof![
            Just(OrfSendRecv::Receive),
            Just(OrfSendRecv::Send),
            Just(OrfSendRecv::Both)
        ],
    )
        .prop_map(|(r#type, send_recv)| OutboundRouteFilterEntry { r#type, send_recv });
    let add_path = prop_oneof![
        Just(AddPathSendRecv::Receive),
        Just(AddPathSendRecv::Send),
        Just(AddPathSendRecv::Both)
    ];
    prop_oneof![
        (afi(), safi()).prop_map(|(afi, safi)| Capability::MultiProtocol { afi, safi }),
        Just(Capability::RouteRefresh),
        vec((afi(), safi(), vec(orf, 0..4)), 0..3).prop_map(|families| {
            Capability::OutboundRouteFilter(
                families
                    .into_iter()
                    .map(|(afi, safi, orfs)| OrfFamily { afi, safi, orfs })
                    .collect(),
            )
        }),
        vec((afi(), safi(), afi()), 1..4).prop_map(|entries| Capability::ExtendedNextHop {
            entries: entries
                .into_iter()
                .map(|(afi, safi, nh_afi)| ExtendedNextHop { afi, safi, nh_afi })
                .collect(),
        }),
        Just(Capability::ExtendedMessage),
        (any::<u8>(), afi()).prop_map(|(header, afi)| Capability::BgpSec {
            header: BgpSecHeader(header),
            afi,
        }),
        vec((afi(), safi(), any::<u8>()), 1..4).prop_map(|entries| {
            Capability::MultiLabelCompat {
                entries: entries
                    .into_iter()
                    .map(|(afi, safi, count)| MultiLabelEntry { afi, safi, count })
                    .collect(),
            }
        }),
        any::<u8>().prop_map(|role| Capability::BgpRole {
            role: BgpRole::from(role),
        }),
        (
            any::<bool>(),
            0u16..4096,
            vec((afi(), safi(), any::<bool>()), 0..4)
        )
            .prop_map(|(restart_state, restart_time, entries)| {
                Capability::GracefulRestart {
                    restart_state,
                    restart_time,
                    entries: entries
                        .into_iter()
                        .map(|(afi, safi, preserve_forwarding)| GracefulRestartEntry {
                            afi,
                            safi,
                            preserve_forwarding,
                        })
                        .collect(),
                }
            }),
        any::<u32>().prop_map(|asn| Capability::FourByteAsn { asn }),
        vec((afi(), safi(), add_path), 1..4).prop_map(|entries| Capability::AddPath {
            entries: entries
                .into_iter()
                .map(|(afi, safi, send_recv)| AddPathEntry {
                    afi,
                    safi,
                    send_recv,
                })
                .collect(),
        }),
        Just(Capability::EnhancedRouteRefresh),
        vec((afi(), safi(), any::<u8>(), 0u32..(1 << 24)), 0..4).prop_map(|entries| {
            Capability::LongLivedGracefulRestart(
                entries
                    .into_iter()
                    .map(|(afi, safi, flags, stale_time)| LongLivedGracefulRestart {
                        afi,
                        safi,
                        flags,
                        stale_time,
                    })
                    .collect(),
            )
        }),
        (
            any::<u8>().prop_filter("a code without a variant", |code| {
                !KNOWN_CAPABILITIES.contains(code)
            }),
            vec(any::<u8>(), 0..20)
        )
            .prop_map(|(type_code, value)| Capability::UnknownCapability { type_code, value }),
    ]
}

fn open() -> impl Strategy<Value = OpenMessage> {
    (
        any::<u16>(),
        any::<u16>(),
        any::<[u8; 4]>(),
        // Enough to go past 255 octets of options now and then, which
        // switches to the extended length of RFC 9072.
        vec(vec(capability(), 0..10), 0..8),
    )
        .prop_map(|(asn, hold_time, identifier, options)| OpenMessage {
            version: 4,
            asn,
            hold_time,
            identifier: Ipv4Addr::from(identifier),
            options: options.into_iter().map(OpenOption::Capabilities).collect(),
        })
}

fn notification() -> impl Strategy<Value = NotificationMessage> {
    // Any code and subcode: a peer may send ones we have no name for.
    (any::<u8>(), any::<u8>(), vec(any::<u8>(), 0..40)).prop_map(|(code, subcode, data)| {
        NotificationMessage {
            error: BgpError::from_code_subcode(code, subcode),
            data,
        }
    })
}

fn message(session: Session) -> impl Strategy<Value = Message> {
    prop_oneof![
        open().prop_map(Message::Open),
        update(session).prop_map(Message::Update),
        notification().prop_map(Message::Notification),
        Just(Message::KeepAlive),
        (afi(), any::<u8>(), safi()).prop_map(|(afi, subtype, safi)| {
            Message::RouteRefresh(RouteRefreshMessage { afi, subtype, safi })
        }),
    ]
}

/// A session, and a well-formed message that fits it.
fn session_and_message() -> impl Strategy<Value = (Session, Message)> {
    session().prop_flat_map(|session| (Just(session), message(session)))
}

// ---------------------------------------------------------------------------
// Properties
// ---------------------------------------------------------------------------

/// Decodes `wire` and checks what must hold for anything the decoder
/// accepts: it can be shown and written, and unless the decoder had to
/// repair it, it reads back as itself.
fn check_decoding(ctx: &ParserContext, wire: &[u8]) -> Result<(), TestCaseError> {
    let Ok(decoded) = Message::decode_with_report(ctx, wire) else {
        return Ok(());
    };
    let (msg, repairs) = (decoded.message, decoded.update_errors);
    // Display of what a peer sent must not be able to panic: it is logged.
    let _ = format!("{msg:?} {repairs:?}");
    if let Message::Update(update) = &msg {
        for prefix in update
            .withdrawn_routes
            .iter()
            .chain(&update.announced_routes)
        {
            let _ = prefix.to_string();
        }
        for attribute in &update.path_attributes {
            match attribute {
                PathAttribute::MpReachNlri(mp) => {
                    mp.prefixes.iter().for_each(|p| drop(p.to_string()))
                }
                PathAttribute::MpUnreachNlri(mp) => {
                    mp.prefixes.iter().for_each(|p| drop(p.to_string()))
                }
                _ => {}
            }
        }
    }

    let mut rewritten = BytesMut::new();
    match msg.to_wire(ctx, &mut rewritten) {
        Ok(()) => {}
        // Withdrawing what a full-size UPDATE announced can need more room
        // than one message has. The result is acted on, never sent.
        Err(EncodeError::TooLong { .. }) if !repairs.is_empty() => return Ok(()),
        Err(e) => {
            return Err(TestCaseError::fail(format!(
                "accepted but cannot be written: {e}"
            )));
        }
    }
    if repairs.is_empty() {
        let again = Message::decode_with_report(ctx, &rewritten)
            .map_err(|e| TestCaseError::fail(format!("written but cannot be read: {e}")))?;
        prop_assert!(again.update_errors.is_empty(), "{:?}", again.update_errors);
        prop_assert_eq!(again.message, msg);
    }
    Ok(())
}

/// Writes a generated message. `None` when it is too big for one BGP
/// message, or some part of it for its length field, which the generators
/// can produce and the encoder must refuse.
fn written(ctx: &ParserContext, msg: &Message) -> Option<BytesMut> {
    let mut wire = BytesMut::new();
    match msg.to_wire(ctx, &mut wire) {
        Ok(()) => Some(wire),
        Err(EncodeError::TooLong { .. }) => None,
    }
}

fn framed(msg_type: u8, body: &[u8]) -> Vec<u8> {
    let mut wire = vec![0xff; 16];
    let length = u16::try_from(usize::from(BGP_HEADER_LEN) + body.len()).unwrap();
    wire.extend_from_slice(&length.to_be_bytes());
    wire.push(msg_type);
    wire.extend_from_slice(body);
    wire
}

/// Ways to damage a message.
#[derive(Debug, Clone)]
enum Damage {
    Set { at: prop::sample::Index, to: u8 },
    Remove { at: prop::sample::Index },
    Insert { at: prop::sample::Index, byte: u8 },
    Truncate { to: prop::sample::Index },
}

fn damage() -> impl Strategy<Value = Damage> {
    prop_oneof![
        (any::<prop::sample::Index>(), any::<u8>()).prop_map(|(at, to)| Damage::Set { at, to }),
        any::<prop::sample::Index>().prop_map(|at| Damage::Remove { at }),
        (any::<prop::sample::Index>(), any::<u8>())
            .prop_map(|(at, byte)| Damage::Insert { at, byte }),
        any::<prop::sample::Index>().prop_map(|to| Damage::Truncate { to }),
    ]
}

/// Applies `damage` to the body of `wire` and makes the header's length
/// right again, so that the damage is the body parsers' problem.
fn damaged(wire: &[u8], damage: &[Damage]) -> Vec<u8> {
    let header = usize::from(BGP_HEADER_LEN);
    let msg_type = wire[header - 1];
    let mut body = wire[header..].to_vec();
    for d in damage {
        match d {
            Damage::Set { at, to } if !body.is_empty() => {
                let i = at.index(body.len());
                body[i] = *to;
            }
            Damage::Remove { at } if !body.is_empty() => {
                body.remove(at.index(body.len()));
            }
            Damage::Insert { at, byte } => body.insert(at.index(body.len() + 1), *byte),
            Damage::Truncate { to } if !body.is_empty() => body.truncate(to.index(body.len())),
            _ => {}
        }
    }
    body.truncate(usize::from(BGP_MAX_MSG_LEN - BGP_HEADER_LEN));
    framed(msg_type, &body)
}

/// 2000 cases a property, which takes about a second. For a longer hunt
/// set `PROPTEST_CASES`, e.g. `PROPTEST_CASES=1000000 cargo test --release
/// -p bgp-packet --test properties`.
fn config() -> ProptestConfig {
    let cases = std::env::var("PROPTEST_CASES")
        .ok()
        .and_then(|cases| cases.parse().ok())
        .unwrap_or(2000);
    ProptestConfig {
        cases,
        // Failures are not saved to a file: one that is found becomes a
        // named test next to the code it broke.
        failure_persistence: None,
        ..ProptestConfig::default()
    }
}

proptest! {
    #![proptest_config(config())]

    #[test]
    fn a_message_that_can_be_written_reads_back_equal((session, msg) in session_and_message()) {
        let ctx = session.ctx();
        let Some(wire) = written(&ctx, &msg) else {
            return Ok(());
        };
        let decoded = Message::decode_with_report(&ctx, &wire)
            .map_err(|e| TestCaseError::fail(format!("{e}")))?;
        // What we write is well formed: nothing for the decoder to repair.
        prop_assert!(decoded.update_errors.is_empty(), "{:?}", decoded.update_errors);
        prop_assert_eq!(&decoded.message, &msg);
    }

    #[test]
    fn noise_is_never_more_than_an_error(session in session(), bytes in vec(any::<u8>(), 0..200)) {
        check_decoding(&session.ctx(), &bytes)?;
    }

    #[test]
    fn noise_behind_a_valid_header_is_never_more_than_an_error(
        session in session(),
        msg_type in 0u8..7,
        body in vec(any::<u8>(), 0..300),
    ) {
        check_decoding(&session.ctx(), &framed(msg_type, &body))?;
    }

    #[test]
    fn a_damaged_message_is_never_more_than_an_error(
        (session, msg) in session_and_message(),
        damage in vec(damage(), 1..4),
    ) {
        let ctx = session.ctx();
        let Some(wire) = written(&ctx, &msg) else {
            return Ok(());
        };
        check_decoding(&ctx, &damaged(&wire, &damage))?;
    }

    /// The parsers below the message level are public too.
    #[test]
    fn noise_is_never_more_than_an_error_for_the_parts(
        session in session(),
        bytes in vec(any::<u8>(), 0..120),
    ) {
        let ctx = session.ctx();
        let _ = OpenMessage::from_wire(&bytes);
        let _ = Capability::from_wire(&bytes);
        let _ = UpdateMessage::decode(&ctx, &bytes);
        let _ = PathAttribute::from_wire(&ctx, &bytes);
    }
}