swanling 0.0.0

use httpmock::{Method::GET, MockRef, MockServer};
use serial_test::serial;

mod common;

use swanling::prelude::*;
use swanling::swanling::SwanlingTaskSet;
use swanling::SwanlingConfiguration;

// Paths used in load tests performed during these tests.
const INDEX_PATH: &str = "/";
const REDIRECT_PATH: &str = "/redirect";
const REDIRECT2_PATH: &str = "/redirect2";
const REDIRECT3_PATH: &str = "/redirect3";
const ABOUT_PATH: &str = "/about.php";

// Indexes to the above paths.
const INDEX_KEY: usize = 0;
const REDIRECT_KEY: usize = 1;
const REDIRECT_KEY2: usize = 2;
const REDIRECT_KEY3: usize = 3;
const ABOUT_KEY: usize = 4;
const SERVER1_INDEX_KEY: usize = 0;
const SERVER1_ABOUT_KEY: usize = 1;
const SERVER1_REDIRECT_KEY: usize = 2;
const SERVER2_INDEX_KEY: usize = 3;
const SERVER2_ABOUT_KEY: usize = 4;

// Load test configuration.
const EXPECT_WORKERS: usize = 4;
const USERS: usize = 9;
const RUN_TIME: usize = 3;

// There are multiple test variations in this file.
#[derive(Clone)]
enum TestType {
    // Chain many different redirects together.
    Chain,
    // Redirect between domains.
    Domain,
    // Permanently redirect between domains.
    Sticky,
}

// Test task.
pub async fn get_index(user: &SwanlingUser) -> SwanlingTaskResult {
    let _swanling = user.get(INDEX_PATH).await?;
    Ok(())
}

// Test task.
pub async fn get_about(user: &SwanlingUser) -> SwanlingTaskResult {
    let _swanling = user.get(ABOUT_PATH).await?;
    Ok(())
}

// Test task.
pub async fn get_redirect(user: &SwanlingUser) -> SwanlingTaskResult {
    // Load REDIRECT_PATH and follow redirects to ABOUT_PATH.
    let mut swanling = user.get(REDIRECT_PATH).await?;

    if let Ok(r) = swanling.response {
        match r.text().await {
            Ok(html) => {
                // Confirm that we followed redirects and loaded the about page.
                if !html.contains("about page") {
                    return user.set_failure(
                        "about page body wrong",
                        &mut swanling.request,
                        None,
                        None,
                    );
                }
            }
            Err(e) => {
                return user.set_failure(
                    format!("unexpected error parsing about page: {}", e).as_str(),
                    &mut swanling.request,
                    None,
                    None,
                );
            }
        }
    }
    Ok(())
}

// Test task.
pub async fn get_domain_redirect(user: &SwanlingUser) -> SwanlingTaskResult {
    let _swanling = user.get(REDIRECT_PATH).await?;
    Ok(())
}

// Sets up the endpoints used to test redirects.
fn setup_mock_server_endpoints<'a>(
    test_type: &TestType,
    server: &'a MockServer,
    server2: Option<&'a MockServer>,
) -> Vec<MockRef<'a>> {
    match test_type {
        TestType::Chain => {
            vec![
                // First set up INDEX_PATH, store in vector at INDEX_KEY.
                server.mock(|when, then| {
                    when.method(GET).path(INDEX_PATH);
                    then.status(200);
                }),
                // Next set up REDIRECT_PATH, store in vector at REDIRECT_KEY.
                server.mock(|when, then| {
                    when.method(GET).path(REDIRECT_PATH);
                    then.status(301).header("Location", REDIRECT2_PATH);
                }),
                // Next set up REDIRECT2_PATH, store in vector at REDIRECT2_KEY.
                server.mock(|when, then| {
                    when.method(GET).path(REDIRECT2_PATH);
                    then.status(302).header("Location", REDIRECT3_PATH);
                }),
                // Next set up REDIRECT3_PATH, store in vector at REDIRECT3_KEY.
                server.mock(|when, then| {
                    when.method(GET).path(REDIRECT3_PATH);
                    then.status(303).header("Location", ABOUT_PATH);
                }),
                // Next set up ABOUT_PATH, store in vector at ABOUT_KEY.
                server.mock(|when, then| {
                    when.method(GET).path(ABOUT_PATH);
                    then.status(200)
                        .body("<HTML><BODY>about page</BODY></HTML>");
                }),
            ]
        }
        TestType::Domain | TestType::Sticky => {
            vec![
                // First set up INDEX_PATH, store in vector at SERVER1_INDEX_KEY.
                server.mock(|when, then| {
                    when.method(GET).path(INDEX_PATH);
                    then.status(200);
                }),
                // Next set up ABOUT_PATH, store in vector at SERVER1_ABOUT_KEY.
                server.mock(|when, then| {
                    when.method(GET).path(ABOUT_PATH);
                    then.status(200)
                        .body("<HTML><BODY>about page</BODY></HTML>");
                }),
                // Next set up REDIRECT_PATH, store in vector at SERVER1_REDIRECT_KEY.
                server.mock(|when, then| {
                    when.method(GET).path(REDIRECT_PATH);
                    then.status(301)
                        .header("Location", &server2.unwrap().url(INDEX_PATH));
                }),
                // Next set up INDEX_PATH on server 2, store in vector at SERVER2_INDEX_KEY.
                server2.unwrap().mock(|when, then| {
                    when.method(GET).path(INDEX_PATH);
                    then.status(200);
                }),
                // Next set up ABOUT_PATH on server 2, store in vector at SERVER2_ABOUT_KEY.
                server2.unwrap().mock(|when, then| {
                    when.method(GET).path(ABOUT_PATH);
                    then.status(200);
                }),
            ]
        }
    }
}

// Build appropriate configuration for these tests.
fn common_build_configuration(
    server: &MockServer,
    sticky: bool,
    worker: Option<bool>,
    manager: Option<usize>,
) -> SwanlingConfiguration {
    if let Some(expect_workers) = manager {
        if sticky {
            common::build_configuration(
                &server,
                vec![
                    "--sticky-follow",
                    "--manager",
                    "--expect-workers",
                    &expect_workers.to_string(),
                    "--users",
                    &USERS.to_string(),
                    "--hatch-rate",
                    &USERS.to_string(),
                    "--run-time",
                    &RUN_TIME.to_string(),
                ],
            )
        } else {
            common::build_configuration(
                &server,
                vec![
                    "--manager",
                    "--expect-workers",
                    &expect_workers.to_string(),
                    "--users",
                    &USERS.to_string(),
                    "--hatch-rate",
                    &USERS.to_string(),
                    "--run-time",
                    &RUN_TIME.to_string(),
                ],
            )
        }
    } else if worker.is_some() {
        common::build_configuration(&server, vec!["--worker"])
    } else if sticky {
        common::build_configuration(
            &server,
            vec![
                "--sticky-follow",
                "--users",
                &USERS.to_string(),
                "--hatch-rate",
                &USERS.to_string(),
                "--run-time",
                &RUN_TIME.to_string(),
            ],
        )
    } else {
        common::build_configuration(
            &server,
            vec![
                "--users",
                &USERS.to_string(),
                "--hatch-rate",
                &USERS.to_string(),
                "--run-time",
                &RUN_TIME.to_string(),
            ],
        )
    }
}

// Helper to confirm all variations generate appropriate results.
fn validate_redirect(test_type: &TestType, mock_endpoints: &[MockRef]) {
    match test_type {
        TestType::Chain => {
            // Confirm that all pages are loaded, even those not requested directly but
            // that are only loaded due to redirects.
            assert!(mock_endpoints[INDEX_KEY].hits() > 0);
            assert!(mock_endpoints[REDIRECT_KEY].hits() > 0);
            assert!(mock_endpoints[REDIRECT_KEY2].hits() > 0);
            assert!(mock_endpoints[REDIRECT_KEY3].hits() > 0);
            assert!(mock_endpoints[ABOUT_KEY].hits() > 0);

            // Confirm the entire redirect chain is loaded the same number of times.
            mock_endpoints[REDIRECT_KEY].assert_hits(mock_endpoints[REDIRECT_KEY2].hits());
            mock_endpoints[REDIRECT_KEY].assert_hits(mock_endpoints[REDIRECT_KEY3].hits());
            mock_endpoints[REDIRECT_KEY].assert_hits(mock_endpoints[ABOUT_KEY].hits());
        }
        TestType::Domain => {
            // All pages on Server1 are loaded.
            assert!(mock_endpoints[SERVER1_INDEX_KEY].hits() > 0);
            assert!(mock_endpoints[SERVER1_REDIRECT_KEY].hits() > 0);
            assert!(mock_endpoints[SERVER1_ABOUT_KEY].hits() > 0);

            // SwanlingUsers are redirected to Server2 correctly.
            assert!(mock_endpoints[SERVER2_INDEX_KEY].hits() > 0);

            // SwanlingUsers do not stick to Server2 and load the other page.
            mock_endpoints[SERVER2_ABOUT_KEY].assert_hits(0);
        }
        TestType::Sticky => {
            // Each SwanlingUser loads the redirect on Server1 one time.
            println!(
                "SERVER1_REDIRECT: {}, USERS: {}",
                mock_endpoints[SERVER1_REDIRECT_KEY].hits(),
                USERS,
            );
            println!(
                "SERVER1_INDEX: {}, SERVER1_ABOUT: {}",
                mock_endpoints[SERVER1_INDEX_KEY].hits(),
                mock_endpoints[SERVER1_ABOUT_KEY].hits(),
            );
            println!(
                "SERVER2_INDEX: {}, SERVER2_ABOUT: {}",
                mock_endpoints[SERVER2_INDEX_KEY].hits(),
                mock_endpoints[SERVER2_ABOUT_KEY].hits(),
            );
            mock_endpoints[SERVER1_REDIRECT_KEY].assert_hits(USERS);

            // Redirected to Server2, no user load anything else on Server1.
            mock_endpoints[SERVER1_INDEX_KEY].assert_hits(0);
            mock_endpoints[SERVER1_ABOUT_KEY].assert_hits(0);

            // All SwanlingUsers go on to load pages on Server2.
            assert!(mock_endpoints[SERVER2_INDEX_KEY].hits() > 0);
            assert!(mock_endpoints[SERVER2_ABOUT_KEY].hits() > 0);
        }
    }
}

// Returns the appropriate taskset needed to build these tests.
fn get_tasks(test_type: &TestType) -> SwanlingTaskSet {
    match test_type {
        TestType::Chain => {
            taskset!("LoadTest")
                // Load index directly.
                .register_task(task!(get_index))
                // Load redirect path, redirect to redirect2 path, redirect to
                // redirect3 path, redirect to about.
                .register_task(task!(get_redirect))
        }
        TestType::Domain | TestType::Sticky => {
            taskset!("LoadTest")
                // First load redirect, takes this request only to another domain.
                .register_task(task!(get_domain_redirect))
                // Load index.
                .register_task(task!(get_index))
                // Load about.
                .register_task(task!(get_about))
        }
    }
}

// Helper to run all standalone tests.
fn run_standalone_test(test_type: TestType) {
    // Start the mock servers.
    let server1 = MockServer::start();
    let server2 = MockServer::start();

    // Setup the endpoints needed for this test on the mock server.
    let mock_endpoints = setup_mock_server_endpoints(&test_type, &server1, Some(&server2));

    // Build appropriate configuration.
    let sticky = match test_type {
        TestType::Sticky => true,
        TestType::Chain | TestType::Domain => false,
    };
    let configuration = common_build_configuration(&server1, sticky, None, None);

    // Run the Swanling Attack.
    common::run_load_test(
        common::build_load_test(configuration, &get_tasks(&test_type), None, None),
        None,
    );

    // Confirm that the load test was actually redirected.
    validate_redirect(&test_type, &mock_endpoints);
}

// Helper to run all standalone tests.
fn run_gaggle_test(test_type: TestType) {
    // Start the mock servers.
    let server1 = MockServer::start();
    let server2 = MockServer::start();

    // Setup the endpoints needed for this test on the mock server.
    let mock_endpoints = setup_mock_server_endpoints(&test_type, &server1, Some(&server2));

    // Build appropriate Worker configuration.
    let sticky = match test_type {
        TestType::Sticky => true,
        TestType::Chain | TestType::Domain => false,
    };
    let worker_configuration = common_build_configuration(&server1, sticky, Some(true), None);

    // Build the load test for the Workers.
    let worker_swanling_attack =
        common::build_load_test(worker_configuration, &get_tasks(&test_type), None, None);

    // Workers launched in own threads, store thread handles.
    let worker_handles = common::launch_gaggle_workers(worker_swanling_attack, EXPECT_WORKERS);

    // Build Manager configuration.
    let manager_configuration =
        common_build_configuration(&server1, sticky, None, Some(EXPECT_WORKERS));

    // Build the load test for the Workers.
    let manager_swanling_attack =
        common::build_load_test(manager_configuration, &get_tasks(&test_type), None, None);

    // Run the Swanling Attack.
    common::run_load_test(manager_swanling_attack, Some(worker_handles));

    // Confirm that the load test was actually redirected.
    validate_redirect(&test_type, &mock_endpoints);
}

#[test]
// Request a page that redirects multiple times with different redirect headers.
fn test_redirect() {
    run_standalone_test(TestType::Chain);
}

#[test]
#[cfg_attr(not(feature = "gaggle"), ignore)]
#[serial]
// Request a page that redirects multiple times with different redirect headers,
// in Regatta mode.
fn test_redirect_gaggle() {
    run_gaggle_test(TestType::Chain);
}

#[test]
// Request a page that redirects to another domain.
// Different domains are simulated with multiple mock servers running on different
// ports.
fn test_domain_redirect() {
    run_standalone_test(TestType::Domain);
}

#[test]
#[cfg_attr(not(feature = "gaggle"), ignore)]
#[serial]
// Request a page that redirects to another domain, in Regatta mode.
// Different domains are simulated with multiple mock servers running on different
// ports.
fn test_domain_redirect_gaggle() {
    run_gaggle_test(TestType::Domain);
}

#[test]
// Request a page that redirects to another domain with --sticky-follow enabled.
// Different domains are simulated with multiple mock servers running on different
// ports.
fn test_sticky_domain_redirect() {
    run_standalone_test(TestType::Sticky);
}

#[test]
#[cfg_attr(not(feature = "gaggle"), ignore)]
#[serial]
// Request a page that redirects to another domain with --sticky-follow enabled, in
// Regatta mode.
// Different domains are simulated with multiple mock servers running on different
// ports.
fn test_sticky_domain_redirect_gaggle() {
    run_gaggle_test(TestType::Sticky);
}