//! The RTR client against a cache played by the test: what it asks for,
//! what it makes of the answers, and what it does when the cache misbehaves
//! or goes away.

use std::net::{Ipv4Addr, SocketAddr};
use std::time::Duration;

use bgp_primitives::ip_prefix::IpPrefix;
use bgp_rpki::rtr::{Pdu, Timers, error_code};
use bgp_rpki::{Algorithm, CacheState, PathValidity, Rpki, RpkiConfig, Update, Validity, Vrp};
use bytes::BytesMut;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::watch;

const TIMEOUT: Duration = Duration::from_secs(5);
const SESSION: u16 = 0x1234;
/// The version of the protocol the client asks for first.
const V: u8 = 2;

fn prefix(s: &str) -> IpPrefix {
    IpPrefix::try_from(s).unwrap()
}

fn vrp(p: &str, max_length: u8, asn: u32) -> Vrp {
    Vrp {
        prefix: prefix(p),
        max_length,
        asn,
    }
}

fn announce(vrp: Vrp) -> Pdu {
    Pdu::Prefix {
        announce: true,
        vrp,
    }
}

fn withdraw(vrp: Vrp) -> Pdu {
    Pdu::Prefix {
        announce: false,
        vrp,
    }
}

fn end_of_data(serial: u32) -> Pdu {
    Pdu::EndOfData {
        session_id: SESSION,
        serial,
        timers: Some(Timers::default()),
    }
}

/// A cache played by the test.
struct Cache {
    listener: TcpListener,
}

/// One connection from the client to the cache.
struct Connection {
    stream: TcpStream,
    buf: BytesMut,
}

impl Cache {
    async fn start() -> Self {
        Self {
            listener: TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).await.unwrap(),
        }
    }

    fn address(&self) -> SocketAddr {
        self.listener.local_addr().unwrap()
    }

    async fn accept(&self) -> Connection {
        let (stream, _) = tokio::time::timeout(TIMEOUT, self.listener.accept())
            .await
            .expect("the client connects")
            .unwrap();
        Connection {
            stream,
            buf: BytesMut::new(),
        }
    }
}

impl Connection {
    /// The next PDU from the client and the version it was sent in.
    async fn received(&mut self) -> (u8, Pdu) {
        let reading = async {
            loop {
                if let Some((version, pdu, used)) = Pdu::decode(&self.buf).unwrap() {
                    let _ = self.buf.split_to(used);
                    return (version, pdu);
                }
                assert_ne!(
                    self.stream.read_buf(&mut self.buf).await.unwrap(),
                    0,
                    "the client closed the connection"
                );
            }
        };
        tokio::time::timeout(TIMEOUT, reading)
            .await
            .expect("the client sends something")
    }

    async fn send(&mut self, version: u8, pdus: &[Pdu]) {
        let mut out = BytesMut::new();
        for pdu in pdus {
            pdu.encode(version, &mut out);
        }
        self.stream.write_all(&out).await.unwrap();
    }

    /// Answers a query with everything: `vrps`, at `serial`.
    async fn send_all(&mut self, vrps: &[Vrp], serial: u32) {
        let mut pdus = vec![Pdu::CacheResponse {
            session_id: SESSION,
        }];
        pdus.extend(vrps.iter().copied().map(announce));
        pdus.push(end_of_data(serial));
        self.send(V, &pdus).await;
    }
}

fn client_of(caches: &[&Cache]) -> (Rpki, watch::Receiver<Update>) {
    let (rpki, _task) = Rpki::spawn(RpkiConfig {
        caches: caches.iter().map(|cache| cache.address()).collect(),
        retry: Some(Duration::from_millis(50)),
        ..RpkiConfig::default()
    });
    let updates = rpki.subscribe();
    (rpki, updates)
}

async fn next_update(updates: &mut watch::Receiver<Update>) -> Update {
    tokio::time::timeout(TIMEOUT, updates.changed())
        .await
        .expect("the VRPs change")
        .unwrap();
    updates.borrow_and_update().clone()
}

/// The ordinary life of a session: everything once, then what changes.
#[tokio::test]
async fn the_client_takes_everything_once_and_then_the_changes() {
    let cache = Cache::start().await;
    let (rpki, mut updates) = client_of(&[&cache]);
    // Nothing is known until a cache has answered.
    assert_eq!(rpki.version(), 0);
    assert!(rpki.current().vrps.is_empty());

    let mut connection = cache.accept().await;
    assert_eq!(connection.received().await, (V, Pdu::ResetQuery));
    let ripe = vrp("193.0.0.0/21", 24, 3333);
    let other = vrp("2001:67c:2e8::/48", 48, 3333);
    connection.send_all(&[ripe, other], 7).await;

    let update = next_update(&mut updates).await;
    assert_eq!(update.version, 1);
    assert_eq!(rpki.version(), 1);
    assert!(
        update.changed.is_none(),
        "the first data can change anything"
    );
    assert_eq!(update.vrps.len(), 2);
    assert_eq!(
        update.vrps.validate(&prefix("193.0.4.0/24"), Some(3333)),
        Validity::Valid
    );
    assert_eq!(
        update.vrps.validate(&prefix("193.0.4.0/24"), Some(15169)),
        Validity::Invalid
    );
    let status = rpki.status();
    assert_eq!(status.state, CacheState::Synchronised);
    assert_eq!((status.session_id, status.serial), (Some(SESSION), Some(7)));
    assert_eq!((status.vrps_ipv4, status.vrps_ipv6), (1, 1));
    assert_eq!(status.cache, Some(cache.address()));

    // The cache has news: the client asks for what changed since serial 7.
    connection
        .send(
            V,
            &[Pdu::SerialNotify {
                session_id: SESSION,
                serial: 8,
            }],
        )
        .await;
    assert_eq!(
        connection.received().await,
        (
            V,
            Pdu::SerialQuery {
                session_id: SESSION,
                serial: 7
            }
        )
    );
    let new = vrp("8.8.8.0/24", 24, 15169);
    connection
        .send(
            V,
            &[
                Pdu::CacheResponse {
                    session_id: SESSION,
                },
                announce(new),
                withdraw(other),
                end_of_data(8),
            ],
        )
        .await;

    let update = next_update(&mut updates).await;
    assert_eq!(update.version, 2);
    assert_eq!(update.vrps.len(), 2);
    assert_eq!(
        update.vrps.validate(&prefix("8.8.8.0/24"), Some(15169)),
        Validity::Valid
    );
    assert_eq!(
        update
            .vrps
            .validate(&prefix("2001:67c:2e8::/48"), Some(3333)),
        Validity::NotFound,
        "withdrawn"
    );
    // Only routes that what came and went covers can have changed.
    let changed = update.changed.expect("a change to what there was");
    assert!(changed.vrps.covers(&prefix("8.8.8.0/24")));
    assert!(changed.vrps.covers(&prefix("2001:67c:2e8::/48")));
    assert!(!changed.vrps.covers(&prefix("193.0.0.0/21")));
    assert_eq!(rpki.status().serial, Some(8));
}

/// A cache that cannot say what changed tells the client to start over.
/// What the client then publishes is the difference to what it had.
#[tokio::test]
async fn a_cache_reset_makes_the_client_ask_for_everything() {
    let cache = Cache::start().await;
    let (rpki, mut updates) = client_of(&[&cache]);
    let mut connection = cache.accept().await;
    connection.received().await;
    let kept = vrp("193.0.0.0/21", 24, 3333);
    let dropped = vrp("8.8.8.0/24", 24, 15169);
    connection.send_all(&[kept, dropped], 7).await;
    next_update(&mut updates).await;

    connection
        .send(
            V,
            &[Pdu::SerialNotify {
                session_id: SESSION,
                serial: 20,
            }],
        )
        .await;
    assert!(matches!(
        connection.received().await.1,
        Pdu::SerialQuery { serial: 7, .. }
    ));
    connection.send(V, &[Pdu::CacheReset]).await;
    assert_eq!(connection.received().await.1, Pdu::ResetQuery);
    let added = vrp("1.1.1.0/24", 24, 13335);
    connection.send_all(&[kept, added], 20).await;

    let update = next_update(&mut updates).await;
    assert_eq!(update.vrps.len(), 2);
    let changed = update.changed.expect("known what changed");
    assert!(changed.vrps.covers(&prefix("8.8.8.0/24")));
    assert!(changed.vrps.covers(&prefix("1.1.1.0/24")));
    assert!(!changed.vrps.covers(&prefix("193.0.0.0/21")), "it stayed");
    assert_eq!(rpki.status().serial, Some(20));
}

/// An answer takes effect whole or not at all: one that contradicts itself
/// changes nothing, and the session is given up.
#[tokio::test]
async fn an_answer_that_contradicts_itself_changes_nothing() {
    let cache = Cache::start().await;
    let (rpki, mut updates) = client_of(&[&cache]);
    let mut connection = cache.accept().await;
    connection.received().await;
    let held = vrp("193.0.0.0/21", 24, 3333);
    connection.send_all(&[held], 7).await;
    next_update(&mut updates).await;

    connection
        .send(
            V,
            &[Pdu::SerialNotify {
                session_id: SESSION,
                serial: 8,
            }],
        )
        .await;
    connection.received().await;
    // A good change, then the withdrawal of something never announced.
    connection
        .send(
            V,
            &[
                Pdu::CacheResponse {
                    session_id: SESSION,
                },
                announce(vrp("1.1.1.0/24", 24, 13335)),
                withdraw(vrp("9.9.9.0/24", 24, 19281)),
                end_of_data(8),
            ],
        )
        .await;

    // The client says what was wrong,
    assert!(matches!(
        connection.received().await,
        (
            V,
            Pdu::ErrorReport {
                code: error_code::WITHDRAWAL_OF_UNKNOWN_RECORD,
                ..
            }
        )
    ));
    // gives the session up and comes back; what it has is what it had,
    // and it asks from where it was.
    let mut again = cache.accept().await;
    assert_eq!(
        again.received().await.1,
        Pdu::SerialQuery {
            session_id: SESSION,
            serial: 7
        }
    );
    let current = rpki.current();
    assert_eq!(current.version, 1);
    assert_eq!(current.vrps.len(), 1);
    assert!(
        rpki.status()
            .last_error
            .is_some_and(|e| e.contains("withdrew a VRP it had not announced"))
    );
}

fn aspa(customer: u32, providers: &[u32]) -> Pdu {
    Pdu::Aspa {
        announce: true,
        customer,
        providers: providers.to_vec(),
    }
}

/// ASPAs come with the VRPs, in version 2. One that is announced again
/// takes the place of the one before.
#[tokio::test]
async fn aspas_are_taken_replaced_and_withdrawn() {
    let cache = Cache::start().await;
    let (rpki, mut updates) = client_of(&[&cache]);
    let mut connection = cache.accept().await;
    assert_eq!(connection.received().await, (V, Pdu::ResetQuery));
    connection
        .send(
            V,
            &[
                Pdu::CacheResponse {
                    session_id: SESSION,
                },
                announce(vrp("193.0.0.0/21", 24, 3333)),
                aspa(3333, &[1101, 3257]),
                aspa(1101, &[0]),
                end_of_data(7),
            ],
        )
        .await;
    let update = next_update(&mut updates).await;
    assert_eq!(update.vrps.len(), 1);
    assert_eq!(update.aspas.len(), 2);
    assert_eq!(update.aspas.providers(3333), Some(&[1101, 3257][..]));
    assert_eq!(rpki.status().aspas, 2);
    assert_eq!(rpki.status().protocol_version, 2);
    assert_eq!(
        update.aspas.verify(&[174, 3333], Algorithm::Upstream),
        PathValidity::Invalid
    );

    connection
        .send(
            V,
            &[Pdu::SerialNotify {
                session_id: SESSION,
                serial: 8,
            }],
        )
        .await;
    connection.received().await;
    connection
        .send(
            V,
            &[
                Pdu::CacheResponse {
                    session_id: SESSION,
                },
                // Another list for 3333, a new ASPA, and one gone.
                aspa(3333, &[174, 1101]),
                aspa(2121, &[3333]),
                Pdu::Aspa {
                    announce: false,
                    customer: 1101,
                    providers: vec![],
                },
                end_of_data(8),
            ],
        )
        .await;
    let update = next_update(&mut updates).await;
    assert_eq!(update.aspas.len(), 2);
    assert_eq!(update.aspas.providers(3333), Some(&[174, 1101][..]));
    assert_eq!(update.aspas.providers(1101), None);
    assert_eq!(
        update.aspas.verify(&[174, 3333], Algorithm::Upstream),
        PathValidity::Valid
    );
    let changed = update.changed.expect("a change to what there was");
    assert_eq!(*changed.aspa_customers, [3333, 2121, 1101].into());
    assert!(changed.vrps.is_empty(), "no VRP came or went");
    assert_eq!(update.vrps.len(), 1);

    // The withdrawal of an ASPA there is not is a contradiction: nothing
    // of the answer is used.
    connection
        .send(
            V,
            &[Pdu::SerialNotify {
                session_id: SESSION,
                serial: 9,
            }],
        )
        .await;
    connection.received().await;
    connection
        .send(
            V,
            &[
                Pdu::CacheResponse {
                    session_id: SESSION,
                },
                aspa(3333, &[6939]),
                Pdu::Aspa {
                    announce: false,
                    customer: 1101,
                    providers: vec![],
                },
                end_of_data(9),
            ],
        )
        .await;
    assert!(matches!(
        connection.received().await.1,
        Pdu::ErrorReport {
            code: error_code::WITHDRAWAL_OF_UNKNOWN_RECORD,
            ..
        }
    ));
    let mut again = cache.accept().await;
    assert!(matches!(
        again.received().await.1,
        Pdu::SerialQuery { serial: 8, .. }
    ));
    let current = rpki.current();
    assert_eq!(current.version, 2);
    assert_eq!(current.aspas.providers(3333), Some(&[174, 1101][..]));
}

/// A cache that speaks version 1 is spoken to in it. It has no ASPAs to
/// give: those that came from a newer cache go when its data takes their
/// place.
#[tokio::test]
async fn a_cache_that_speaks_version_one_has_no_aspas() {
    let cache = Cache::start().await;
    let (rpki, mut updates) = client_of(&[&cache]);
    let mut connection = cache.accept().await;
    connection.received().await;
    let ripe = vrp("193.0.0.0/21", 24, 3333);
    connection
        .send(
            V,
            &[
                Pdu::CacheResponse {
                    session_id: SESSION,
                },
                announce(ripe),
                aspa(3333, &[1101]),
                end_of_data(7),
            ],
        )
        .await;
    assert_eq!(next_update(&mut updates).await.aspas.len(), 1);

    // The validator is replaced by an older one at the same address.
    drop(connection);
    let mut connection = cache.accept().await;
    assert!(matches!(
        connection.received().await,
        (V, Pdu::SerialQuery { serial: 7, .. })
    ));
    connection
        .send(
            1,
            &[Pdu::ErrorReport {
                code: error_code::UNSUPPORTED_VERSION,
                text: "version 1 at most".into(),
            }],
        )
        .await;
    // Asked again in version 1, and for everything: the serial was of
    // another version's data.
    let mut connection = cache.accept().await;
    assert_eq!(connection.received().await, (1, Pdu::ResetQuery));
    connection
        .send(
            1,
            &[
                Pdu::CacheResponse {
                    session_id: SESSION,
                },
                announce(ripe),
                end_of_data(3),
            ],
        )
        .await;
    let update = next_update(&mut updates).await;
    assert!(update.aspas.is_empty());
    assert_eq!(update.vrps.len(), 1);
    let changed = update.changed.expect("the difference to what there was");
    assert_eq!(*changed.aspa_customers, [3333].into());
    assert!(changed.vrps.is_empty());
    let status = rpki.status();
    assert_eq!((status.protocol_version, status.aspas), (1, 0));

    // An ASPA in a version that has none is not RTR.
    connection
        .stream
        .write_all(&[1, 11, 1, 0, 0, 0, 0, 16, 0, 0, 0, 1, 0, 0, 0, 2])
        .await
        .unwrap();
    assert!(matches!(
        connection.received().await,
        (
            1,
            Pdu::ErrorReport {
                code: error_code::UNSUPPORTED_PDU_TYPE,
                ..
            }
        )
    ));
}

/// A validator that only speaks version 0 says so, by an error or by
/// answering in version 0. The client asks again in version 0.
#[tokio::test]
async fn a_cache_that_speaks_version_zero_is_spoken_to_in_it() {
    for refuses_with_an_error in [true, false] {
        let cache = Cache::start().await;
        let (rpki, mut updates) = client_of(&[&cache]);
        let mut connection = cache.accept().await;
        assert_eq!(connection.received().await, (V, Pdu::ResetQuery));
        if refuses_with_an_error {
            connection
                .send(
                    0,
                    &[Pdu::ErrorReport {
                        code: error_code::UNSUPPORTED_VERSION,
                        text: "version 0 only".into(),
                    }],
                )
                .await;
        } else {
            connection
                .send(
                    0,
                    &[Pdu::CacheResponse {
                        session_id: SESSION,
                    }],
                )
                .await;
        }

        let mut connection = cache.accept().await;
        assert_eq!(connection.received().await, (0, Pdu::ResetQuery));
        connection
            .send(
                0,
                &[
                    Pdu::CacheResponse {
                        session_id: SESSION,
                    },
                    announce(vrp("193.0.0.0/21", 24, 3333)),
                    Pdu::EndOfData {
                        session_id: SESSION,
                        serial: 3,
                        timers: None,
                    },
                ],
            )
            .await;
        assert_eq!(next_update(&mut updates).await.vrps.len(), 1);
        assert_eq!(rpki.status().protocol_version, 0);
    }
}

/// When the cache goes away the data is kept and used while the client
/// tries again, and the next cache is tried. A cache the data did not come
/// from is asked for everything.
#[tokio::test]
async fn a_lost_cache_is_replaced_by_the_next_and_the_data_is_kept() {
    let first = Cache::start().await;
    let second = Cache::start().await;
    let (rpki, mut updates) = client_of(&[&first, &second]);

    let mut connection = first.accept().await;
    connection.received().await;
    connection
        .send_all(&[vrp("193.0.0.0/21", 24, 3333)], 7)
        .await;
    next_update(&mut updates).await;

    drop(connection);
    drop(first);
    let mut connection = second.accept().await;
    assert_eq!(connection.received().await.1, Pdu::ResetQuery);
    // Still in use meanwhile.
    assert_eq!(rpki.current().vrps.len(), 1);

    connection
        .send_all(
            &[vrp("193.0.0.0/21", 24, 3333), vrp("1.1.1.0/24", 24, 13335)],
            100,
        )
        .await;
    let update = next_update(&mut updates).await;
    assert_eq!(update.vrps.len(), 2);
    let changed = update.changed.expect("the difference to what there was");
    assert!(changed.vrps.covers(&prefix("1.1.1.0/24")));
    assert!(!changed.vrps.covers(&prefix("193.0.0.0/21")));
    assert_eq!(rpki.status().cache, Some(second.address()));
}

/// Data that no cache has confirmed for the expire interval is dropped:
/// routes are then neither valid nor invalid.
#[tokio::test]
async fn data_that_is_not_confirmed_in_time_is_dropped() {
    let cache = Cache::start().await;
    let (rpki, _task) = Rpki::spawn(RpkiConfig {
        caches: vec![cache.address()],
        // Data never expires sooner than a refresh and a retry take.
        refresh: Some(Duration::from_millis(200)),
        retry: Some(Duration::from_millis(50)),
        expire: Some(Duration::from_millis(300)),
    });
    let mut updates = rpki.subscribe();
    let mut connection = cache.accept().await;
    connection.received().await;
    connection
        .send_all(&[vrp("193.0.0.0/21", 24, 3333)], 7)
        .await;
    assert_eq!(next_update(&mut updates).await.vrps.len(), 1);

    // The cache goes and does not come back.
    drop(connection);
    drop(cache);
    let update = next_update(&mut updates).await;
    assert!(update.vrps.is_empty());
    assert!(update.changed.is_none(), "anything may have changed");
    assert_eq!(
        update.vrps.validate(&prefix("193.0.0.0/21"), Some(15169)),
        Validity::NotFound
    );
    let status = rpki.status();
    assert_eq!(status.state, CacheState::Connecting);
    assert_eq!(status.serial, None);
}

/// A cache that reports an error, or sends what is not RTR, is given up
/// on and tried again; nothing of what it sent is used.
#[tokio::test]
async fn a_cache_that_misbehaves_is_given_up_on() {
    let cache = Cache::start().await;
    let (rpki, _updates) = client_of(&[&cache]);

    let mut connection = cache.accept().await;
    connection.received().await;
    connection
        .send(
            V,
            &[Pdu::ErrorReport {
                code: error_code::NO_DATA_AVAILABLE,
                text: "still validating".into(),
            }],
        )
        .await;
    let mut connection = cache.accept().await;
    assert_eq!(connection.received().await.1, Pdu::ResetQuery);
    assert!(
        rpki.status()
            .last_error
            .is_some_and(|e| e.contains("still validating"))
    );

    // Half an answer, then noise.
    connection
        .send(
            V,
            &[
                Pdu::CacheResponse {
                    session_id: SESSION,
                },
                announce(vrp("193.0.0.0/21", 24, 3333)),
            ],
        )
        .await;
    connection
        .stream
        .write_all(&[1, 99, 0, 0, 0, 0, 0, 8])
        .await
        .unwrap();
    let mut connection = cache.accept().await;
    assert_eq!(connection.received().await.1, Pdu::ResetQuery);
    assert!(rpki.current().vrps.is_empty());
    assert_eq!(rpki.version(), 0);
}