irtt-rs/crates/irtt-client/tests/managed_authentication.rs

#![cfg(feature = "tokio")]

use std::{collections::HashSet, time::Duration};

use irtt_client::{managed::*, Authentication, ClientConfig, ClientEvent, HmacKey};
use irtt_server::{Server, ServerConfig};

fn config(auth: Authentication) -> ManagedClientConfig {
    ManagedClientConfig {
        client: ClientConfig {
            auth,
            duration: None,
            interval: Duration::from_millis(20),
            probe_timeout: Duration::from_millis(100),
            open_timeouts: vec![Duration::from_millis(200)],
            ..ClientConfig::default()
        },
        completion: ManagedCompletionPolicy::ExplicitStop,
        final_drain: Duration::ZERO,
        ..ManagedClientConfig::default()
    }
}

fn runtime() -> tokio::runtime::Runtime {
    tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap()
}

#[test]
fn managed_authentication_inherits_disables_and_overrides_on_creation_and_update() {
    runtime().block_on(async {
        tokio::time::timeout(Duration::from_secs(5), async {
            let shared = Authentication::Hmac(HmacKey::new(b"shared-test-key".as_slice()));
            let replacement =
                Authentication::Hmac(HmacKey::new(b"replacement-test-key".as_slice()));
            let empty = Authentication::Hmac(HmacKey::new(Vec::new()));
            let mut targets = Vec::new();
            let mut servers = Vec::new();
            for (id, key, auth) in [
                (
                    "inherit",
                    Some(b"shared-test-key".as_slice()),
                    TargetAuth::Inherit,
                ),
                (
                    "disable",
                    None,
                    TargetAuth::Override(Authentication::Unauthenticated),
                ),
                (
                    "replace",
                    Some(b"replacement-test-key".as_slice()),
                    TargetAuth::Override(replacement),
                ),
                ("empty", Some(&[]), TargetAuth::Override(empty)),
            ] {
                let server_config = key.map_or_else(ServerConfig::default, |key| {
                    ServerConfig::default().with_hmac_key(key)
                });
                let mut server = Server::bind("127.0.0.1:0".parse().unwrap(), server_config)
                    .await
                    .unwrap();
                let mut target =
                    ManagedTargetConfig::new(id, server.local_addr().unwrap().to_string());
                target.auth = auth;
                targets.push(target);
                servers.push(tokio::spawn(async move {
                    server.run(std::future::pending()).await
                }));
            }
            for initial in [true, false] {
                let (task, handle) = ManagedClient::task(
                    config(shared.clone()),
                    if initial { targets.clone() } else { vec![] },
                )
                .unwrap();
                let mut events = handle.subscribe().unwrap();
                let run = tokio::spawn(task);
                if !initial {
                    handle
                        .update_targets(targets.clone())
                        .unwrap()
                        .await
                        .unwrap();
                }
                let mut replied = HashSet::new();
                while replied.len() < targets.len() {
                    if let ManagedEvent::Client {
                        target,
                        event: ClientEvent::EchoReply { .. },
                    } = events.recv().await.unwrap()
                    {
                        replied.insert(target.id);
                    }
                }
                handle.stop().await;
                let outcome = run.await.unwrap();
                assert_eq!(outcome.end_reason, ManagedEndReason::StopRequested);
                assert_eq!(outcome.successful_target_outcomes, targets.len() as u64);
                assert!(outcome.recent_target_outcomes.iter().all(|target| {
                    target.replies_received > 0 && target.cleanup_failure.is_none()
                }));
            }
            for server in servers {
                server.abort();
                assert!(server.await.unwrap_err().is_cancelled());
            }
        })
        .await
        .expect("every managed authentication mode must receive replies");
    });
}

#[test]
fn changing_auth_specification_replaces_generation_even_when_resolved_auth_is_equal() {
    runtime().block_on(async {
        tokio::time::timeout(Duration::from_secs(5), async {
            for auth in [
                Authentication::Unauthenticated,
                Authentication::Hmac(HmacKey::new(b"shared-test-key".as_slice())),
            ] {
                let server_config = match &auth {
                    Authentication::Unauthenticated => ServerConfig::default(),
                    Authentication::Hmac(key) => {
                        ServerConfig::default().with_hmac_key(key.as_bytes())
                    }
                };
                let mut server = Server::bind("127.0.0.1:0".parse().unwrap(), server_config)
                    .await
                    .unwrap();
                let mut configured =
                    ManagedTargetConfig::new("one", server.local_addr().unwrap().to_string());
                let service = tokio::spawn(async move { server.run(std::future::pending()).await });
                let (task, handle) =
                    ManagedClient::task(config(auth.clone()), vec![configured.clone()]).unwrap();
                let mut events = handle.subscribe().unwrap();
                let run = tokio::spawn(task);
                let mut current = handle.status().targets[0].target.clone();
                // Wait for a real reply before changing the running session.
                loop {
                    if matches!(events.recv().await.unwrap(), ManagedEvent::Client {
                        target, event: ClientEvent::EchoReply { .. }
                    } if target == current)
                    {
                        break;
                    }
                }
                for specification in [TargetAuth::Override(auth.clone()), TargetAuth::Inherit] {
                    configured.auth = specification;
                    let applied = handle
                        .update_targets(vec![configured.clone()])
                        .unwrap()
                        .await
                        .unwrap();
                    let next = applied
                        .status
                        .targets
                        .iter()
                        .find(|target| target.desired)
                        .unwrap()
                        .target
                        .clone();
                    assert!(next.generation > current.generation);
                    let mut retired = false;
                    let mut replied = false;
                    while !retired || !replied {
                        match events.recv().await.unwrap() {
                            ManagedEvent::TargetFinished { outcome }
                                if outcome.target == current =>
                            {
                                assert_eq!(outcome.end_reason, ManagedTargetEndReason::Replaced);
                                assert!(outcome.cleanup_failure.is_none());
                                retired = true;
                            }
                            ManagedEvent::Client {
                                target,
                                event: ClientEvent::EchoReply { .. },
                            } if target == next => {
                                replied = true;
                            }
                            _ => {}
                        }
                    }
                    let identical = handle
                        .update_targets(vec![configured.clone()])
                        .unwrap()
                        .await
                        .unwrap();
                    assert_eq!(
                        identical
                            .status
                            .targets
                            .iter()
                            .find(|target| target.desired)
                            .unwrap()
                            .target,
                        next
                    );
                    current = next;
                }
                handle.stop().await;
                assert_eq!(run.await.unwrap().successful_target_outcomes, 3);
                service.abort();
                assert!(service.await.unwrap_err().is_cancelled());
            }
        })
        .await
        .expect("auth specification changes must replace and reopen target generations");
    });
}