🐸 Frogue(Fed)

~ololduck/Frogue

main / frogue / tests / ssh_throttling.rs

Raw

//! SSH connection throttling integration tests.
//!
//! These tests verify that connection limits and timeouts work correctly.
//!
//! TODO: write an helper to check if stream is closed.
//! TODO: still need to write more accurate tests, but my guess is that we would
//! need to hook into russh. Maybe improve the tests with some timers?

mod support;

use std::net::TcpStream;
use std::time::Duration;

use support::{Server, free_port};

fn spawn_with_env(env: &[(&str, &str)]) -> (Server, u16) {
    let port = free_port();
    let mut server = Server::spawn_with_env(&format!("127.0.0.1:{port}"), env);
    server.wait_for_connection(port);
    let ssh_port = server.ssh_port();
    server.wait_for_connection(ssh_port);
    (server, port)
}

/// A simple SSH client that just connects and waits.
fn connect_ssh(ssh_port: u16) -> TcpStream {
    TcpStream::connect(format!("127.0.0.1:{ssh_port}")).expect("should connect to SSH port")
}

#[test]
fn tui_timeout_disconnects_idle_guest() {
    // Set TUI timeout to 1 second
    let env = [("FROGUE_SSH_TUI_TIMEOUT_SECONDS", "1")];
    let (server, _) = spawn_with_env(&env);
    let ssh_port = server.ssh_port();

    // Connect and let it sit idle
    let _stream = connect_ssh(ssh_port);

    // Wait for timeout + buffer
    std::thread::sleep(Duration::from_secs(2));

    // The connection should have been closed by the server
    // We can't easily test this without implementing SSH protocol,
    // but we can at least verify the server is still running
    let _stream2 = connect_ssh(ssh_port);
}

#[test]
fn guest_idle_timeout_disconnects_guest() {
    // Set guest idle timeout to 1 second
    let env = [("FROGUE_SSH_GUEST_IDLE_TIMEOUT_SECONDS", "1")];
    let (server, _) = spawn_with_env(&env);
    let ssh_port = server.ssh_port();

    // Connect as a guest (unregistered key)
    let _stream = connect_ssh(ssh_port);

    // Wait for timeout + buffer
    std::thread::sleep(Duration::from_secs(2));

    // Server should still be running
    let _stream2 = connect_ssh(ssh_port);
}

#[test]
fn max_guests_limit_rejects_excess() {
    // Allow only 1 guest connection
    let env = [("FROGUE_SSH_MAX_GUEST_CONNECTIONS", "1")];
    let (server, _) = spawn_with_env(&env);
    let ssh_port = server.ssh_port();

    // First guest should connect successfully
    let stream1 = connect_ssh(ssh_port);

    // Second guest should be rejected
    let _stream2 = connect_ssh(ssh_port);
    drop(stream1);
    // Server should still be running
    let _stream3 = connect_ssh(ssh_port);
}

#[test]
fn max_connections_limit_closes_excess() {
    // Allow only 1 total connection
    let env = [("FROGUE_SSH_MAX_CONNECTIONS", "1")];
    let (server, _) = spawn_with_env(&env);
    let ssh_port = server.ssh_port();

    // First connection should succeed
    let stream1 = connect_ssh(ssh_port);

    // Second connection should be closed immediately
    let _stream2 = connect_ssh(ssh_port);
    drop(stream1);

    // Server should still be running
    let _stream3 = connect_ssh(ssh_port);
}