Part 4: Async Programming with Tokio - Concurrent Execution

Introduction

Testing 280 payloads sequentially would take minutes. My WAF scanner needed to test multiple payloads concurrently while respecting rate limits. Rust's async/await with Tokio made this possible—testing dozens of requests simultaneously while maintaining memory safety.

This article shows how I built concurrent scanning into simple-waf-scanner using async Rust and the Tokio runtime.

Why Async in Rust?

The Problem: Waiting on I/O

// Synchronous code - blocking
fn test_payloads_sync(payloads: &[Payload]) {
    for payload in payloads {
        let response = make_http_request(payload);  // Blocks here
        process_response(response);
    }
}
// 280 payloads × 100ms per request = 28 seconds!

The Solution: Async/Await

// Asynchronous code - concurrent
async fn test_payloads_async(payloads: &[Payload]) {
    let mut tasks = Vec::new();
    
    for payload in payloads {
        let task = tokio::spawn(async move {
            let response = make_http_request(payload).await;  // Doesn't block
            process_response(response).await
        });
        tasks.push(task);
    }
    
    for task in tasks {
        task.await.unwrap();
    }
}
// 280 payloads running concurrently = ~2 seconds!

Setting Up Tokio

Add to Cargo.toml

[dependencies]
tokio = { version = "1", features = ["full"] }
reqwest = { version = "0.11", features = ["json"] }
futures = "0.3"

Async Main Function

// The #[tokio::main] macro sets up the async runtime
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("WAF Scanner starting...");
    
    let config = Config::load()?;
    let results = perform_scan(&config).await?;
    
    println!("Scan complete: {} findings", results.len());
    Ok(())
}

Async/Await Basics

Async Functions Return Futures

// Regular function
fn get_data() -> String {
    "data".to_string()
}

// Async function - returns Future<Output = String>
async fn get_data_async() -> String {
    "data".to_string()
}

// Must await the Future to get the value
#[tokio::main]
async fn main() {
    let data = get_data_async().await;  // .await executes the Future
    println!("{}", data);
}

Futures Are Lazy

async fn expensive_operation() -> i32 {
    println!("Running expensive operation");
    42
}

#[tokio::main]
async fn main() {
    // Creating the future doesn't run it
    let future = expensive_operation();
    println!("Future created");  // Prints first
    
    // .await actually runs it
    let result = future.await;   // "Running expensive operation" prints here
    println!("Result: {}", result);
}

Real Example: HTTP Requests

Single Async Request

use reqwest::Client;

async fn fetch_url(client: &Client, url: &str) -> Result<String, reqwest::Error> {
    let response = client
        .get(url)
        .send()
        .await?;  // Yields control while waiting for response
    
    let body = response.text().await?;  // Yields while reading body
    Ok(body)
}

#[tokio::main]
async fn main() {
    let client = Client::new();
    
    match fetch_url(&client, "https://example.com").await {
        Ok(body) => println!("Response: {} bytes", body.len()),
        Err(e) => eprintln!("Error: {}", e),
    }
}

Concurrent Requests

use tokio;
use std::time::Duration;

async fn test_payload(
    client: &Client,
    url: &str,
    payload: &str,
) -> Result<Response, Error> {
    let response = client
        .get(url)
        .query(&[("test", payload)])
        .timeout(Duration::from_secs(10))
        .send()
        .await?;
    
    Ok(Response {
        status: response.status().as_u16(),
        body: response.text().await?,
    })
}

async fn test_multiple_payloads(
    client: &Client,
    url: &str,
    payloads: Vec<String>,
) {
    let mut tasks = Vec::new();
    
    for payload in payloads {
        // Clone what needs to be moved into task
        let client = client.clone();
        let url = url.to_string();
        
        // Spawn concurrent task
        let task = tokio::spawn(async move {
            test_payload(&client, &url, &payload).await
        });
        
        tasks.push(task);
    }
    
    // Wait for all tasks to complete
    for task in tasks {
        match task.await {
            Ok(Ok(response)) => {
                println!("Status: {}", response.status);
            }
            Ok(Err(e)) => {
                eprintln!("Request error: {}", e);
            }
            Err(e) => {
                eprintln!("Task error: {}", e);
            }
        }
    }
}

Controlling Concurrency

Problem: Too Many Concurrent Requests

// Bad: Spawns 280 tasks at once!
for payload in &all_payloads {  // 280 payloads
    tokio::spawn(test_payload(payload));
}
// Overwhelms server, wastes resources

Solution: Semaphore for Limiting

use tokio::sync::Semaphore;
use std::sync::Arc;

async fn test_payloads_limited(
    client: Arc<Client>,
    payloads: Vec<Payload>,
    max_concurrent: usize,
) {
    // Semaphore limits concurrent tasks
    let semaphore = Arc::new(Semaphore::new(max_concurrent));
    let mut tasks = Vec::new();
    
    for payload in payloads {
        let permit = semaphore.clone().acquire_owned().await.unwrap();
        let client = client.clone();
        
        let task = tokio::spawn(async move {
            let result = test_payload(&client, &payload).await;
            drop(permit);  // Release permit when done
            result
        });
        
        tasks.push(task);
    }
    
    for task in tasks {
        task.await.unwrap();
    }
}

Better Solution: Stream with Concurrency Limit

use futures::stream::{self, StreamExt};

async fn test_payloads_stream(
    client: &Client,
    payloads: Vec<Payload>,
    concurrency: usize,
) -> Vec<Finding> {
    stream::iter(payloads)
        .map(|payload| async move {
            test_payload(client, &payload).await
        })
        .buffer_unordered(concurrency)  // Run `concurrency` at a time
        .filter_map(|result| async move {
            match result {
                Ok(Some(finding)) => Some(finding),
                Ok(None) => None,
                Err(e) => {
                    eprintln!("Error: {}", e);
                    None
                }
            }
        })
        .collect()
        .await
}

Rate Limiting

Adding Delays Between Requests

use tokio::time::{sleep, Duration};

async fn test_with_rate_limit(
    client: &Client,
    payloads: &[Payload],
    delay_ms: u64,
) {
    for payload in payloads {
        let result = test_payload(client, payload).await;
        println!("Tested: {}", payload.id);
        
        // Wait before next request
        sleep(Duration::from_millis(delay_ms)).await;
    }
}

Token Bucket Rate Limiter

use governor::{Quota, RateLimiter};
use std::num::NonZeroU32;

async fn test_with_governor(
    client: &Client,
    payloads: Vec<Payload>,
) {
    // Allow 10 requests per second
    let limiter = RateLimiter::direct(
        Quota::per_second(NonZeroU32::new(10).unwrap())
    );
    
    for payload in payloads {
        // Wait until we can make a request
        limiter.until_ready().await;
        
        test_payload(client, &payload).await;
    }
}

Real Implementation from My Scanner

Complete concurrent scanning with limits:

use std::sync::Arc;
use tokio::sync::Semaphore;
use futures::stream::{self, StreamExt};

pub struct Scanner {
    client: Arc<Client>,
    config: Config,
}

pub struct Config {
    pub target_url: String,
    pub concurrency: usize,
    pub delay_ms: u64,
    pub timeout_secs: u64,
}

impl Scanner {
    pub fn new(config: Config) -> Self {
        let client = Client::builder()
            .timeout(Duration::from_secs(config.timeout_secs))
            .build()
            .unwrap();
        
        Scanner {
            client: Arc::new(client),
            config,
        }
    }
    
    pub async fn scan(&self, payloads: Vec<Payload>) -> Vec<Finding> {
        let semaphore = Arc::new(Semaphore::new(self.config.concurrency));
        
        stream::iter(payloads)
            .map(|payload| {
                let client = self.client.clone();
                let url = self.config.target_url.clone();
                let delay = self.config.delay_ms;
                let permit = semaphore.clone();
                
                async move {
                    // Acquire permit (limits concurrency)
                    let _permit = permit.acquire().await.unwrap();
                    
                    // Test payload
                    let result = self.test_single_payload(
                        &client,
                        &url,
                        &payload
                    ).await;
                    
                    // Rate limiting delay
                    if delay > 0 {
                        tokio::time::sleep(Duration::from_millis(delay)).await;
                    }
                    
                    result
                }
            })
            .buffer_unordered(self.config.concurrency)
            .filter_map(|result| async {
                result.ok().flatten()
            })
            .collect()
            .await
    }
    
    async fn test_single_payload(
        &self,
        client: &Client,
        url: &str,
        payload: &Payload,
    ) -> Result<Option<Finding>> {
        let response = client
            .get(url)
            .query(&[("q", &payload.value)])
            .send()
            .await?;
        
        if self.is_bypass(&response) {
            Ok(Some(Finding::new(payload, response)))
        } else {
            Ok(None)
        }
    }
    
    fn is_bypass(&self, response: &Response) -> bool {
        response.status().is_success()
    }
}

Async Error Handling

Propagating Errors in Async

async fn load_and_process() -> Result<ProcessedData, Error> {
    let data = load_data().await?;  // Async operation
    let validated = validate_data(&data)?;  // Sync operation
    let processed = process_data(validated).await?;  // Async operation
    Ok(processed)
}

Timeout for Async Operations

use tokio::time::{timeout, Duration};

async fn with_timeout() -> Result<Response, Error> {
    match timeout(
        Duration::from_secs(5),
        make_request()
    ).await {
        Ok(Ok(response)) => Ok(response),
        Ok(Err(e)) => Err(e),
        Err(_) => Err(Error::Timeout),
    }
}

Channels for Communication

Sending Results Between Tasks

use tokio::sync::mpsc;

async fn scan_with_progress(payloads: Vec<Payload>) {
    let (tx, mut rx) = mpsc::channel(100);
    
    // Spawn scanning tasks
    for payload in payloads {
        let tx = tx.clone();
        tokio::spawn(async move {
            let result = test_payload(&payload).await;
            tx.send(result).await.unwrap();
        });
    }
    drop(tx);  // Close sender
    
    // Receive and display results
    let mut findings = Vec::new();
    while let Some(result) = rx.recv().await {
        if let Ok(Some(finding)) = result {
            println!("Found: {}", finding.payload_id);
            findings.push(finding);
        }
    }
    
    println!("Total findings: {}", findings.len());
}

Select - Racing Futures

use tokio::select;

async fn first_to_respond(urls: Vec<String>) -> Option<String> {
    let mut tasks = Vec::new();
    
    for url in urls {
        tasks.push(fetch_url(&url));
    }
    
    // Return first successful response
    loop {
        select! {
            result = tasks[0] => return result.ok(),
            result = tasks[1] => return result.ok(),
            result = tasks[2] => return result.ok(),
        }
    }
}

Join - Waiting for Multiple Futures

use tokio::join;

async fn parallel_operations() -> (Result<A>, Result<B>, Result<C>) {
    // Run all three concurrently, wait for all
    let (a, b, c) = join!(
        operation_a(),
        operation_b(),
        operation_c()
    );
    
    (a, b, c)
}

// Or with try_join - stop on first error
use tokio::try_join;

async fn parallel_fallible() -> Result<(A, B, C)> {
    let (a, b, c) = try_join!(
        operation_a(),
        operation_b(),
        operation_c()
    )?;
    
    Ok((a, b, c))
}

Testing Async Code

#[cfg(test)]
mod tests {
    use super::*;
    
    #[tokio::test]
    async fn test_concurrent_scanning() {
        let client = Arc::new(Client::new());
        let payloads = vec![
            Payload::new("test1"),
            Payload::new("test2"),
        ];
        
        let results = test_payloads(&client, payloads).await;
        assert_eq!(results.len(), 2);
    }
    
    #[tokio::test]
    async fn test_timeout() {
        let result = timeout(
            Duration::from_millis(100),
            slow_operation()
        ).await;
        
        assert!(result.is_err());  // Should timeout
    }
}

Key Takeaways

1. Async/await for I/O: Don't block on network/disk operations

2. Tokio runtime: Executes async tasks concurrently

3. Spawning tasks: tokio::spawn for concurrent execution

4. Limit concurrency: Use semaphores or streams

5. Rate limiting: Add delays between requests

6. Channels: Communicate between tasks

7. Error handling: Async errors work like sync errors

8. Testing: Use #[tokio::test] for async tests

Common Patterns

// Pattern 1: Spawn and forget
tokio::spawn(async {
    background_task().await;
});

// Pattern 2: Spawn and collect
let tasks: Vec<_> = items.into_iter()
    .map(|item| tokio::spawn(process(item)))
    .collect();
for task in tasks {
    task.await.unwrap();
}

// Pattern 3: Stream with concurrency
stream::iter(items)
    .map(|item| async { process(item).await })
    .buffer_unordered(concurrency)
    .collect::<Vec<_>>()
    .await

// Pattern 4: Select first
select! {
    result = future1 => handle_1(result),
    result = future2 => handle_2(result),
}

Next in Series

In Part 5: HTTP Clients and Real-World Integration, we'll put everything together—building the complete WAF scanner with HTTP requests, JSON handling, CLI parsing, and publishing to crates.io.

---

Based on implementing concurrent payload testing in simple-waf-scanner, achieving 10-50x speedup over sequential execution.