smartproxy/rust/crates/rustproxy/src/challenge_server.rs

//! HTTP-01 ACME challenge server.
//!
//! A lightweight HTTP server that serves ACME challenge responses at
//! `/.well-known/acme-challenge/<token>`.

use std::sync::Arc;

use bytes::Bytes;
use dashmap::DashMap;
use http_body_util::Full;
use hyper::body::Incoming;
use hyper::{Request, Response, StatusCode};
use hyper_util::rt::TokioIo;
use tokio::net::TcpListener;
use tokio_util::sync::CancellationToken;
use tracing::{debug, info, error};

/// ACME HTTP-01 challenge server.
pub struct ChallengeServer {
    /// Token -> key authorization mapping
    challenges: Arc<DashMap<String, String>>,
    /// Cancellation token to stop the server
    cancel: CancellationToken,
    /// Server task handle
    handle: Option<tokio::task::JoinHandle<()>>,
}

impl ChallengeServer {
    /// Create a new challenge server (not yet started).
    pub fn new() -> Self {
        Self {
            challenges: Arc::new(DashMap::new()),
            cancel: CancellationToken::new(),
            handle: None,
        }
    }

    /// Register a challenge token -> key_authorization mapping.
    pub fn set_challenge(&self, token: String, key_authorization: String) {
        debug!("Registered ACME challenge: token={}", token);
        self.challenges.insert(token, key_authorization);
    }

    /// Remove a challenge token.
    pub fn remove_challenge(&self, token: &str) {
        self.challenges.remove(token);
    }

    /// Start the challenge server on the given port.
    pub async fn start(&mut self, port: u16) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let addr = format!("0.0.0.0:{}", port);
        let listener = TcpListener::bind(&addr).await?;
        info!("ACME challenge server listening on port {}", port);

        let challenges = Arc::clone(&self.challenges);
        let cancel = self.cancel.clone();

        let handle = tokio::spawn(async move {
            loop {
                tokio::select! {
                    _ = cancel.cancelled() => {
                        info!("ACME challenge server stopping");
                        break;
                    }
                    result = listener.accept() => {
                        match result {
                            Ok((stream, _)) => {
                                let challenges = Arc::clone(&challenges);
                                tokio::spawn(async move {
                                    let io = TokioIo::new(stream);
                                    let service = hyper::service::service_fn(move |req: Request<Incoming>| {
                                        let challenges = Arc::clone(&challenges);
                                        async move {
                                            Self::handle_request(req, &challenges)
                                        }
                                    });

                                    let conn = hyper::server::conn::http1::Builder::new()
                                        .serve_connection(io, service);

                                    if let Err(e) = conn.await {
                                        debug!("Challenge server connection error: {}", e);
                                    }
                                });
                            }
                            Err(e) => {
                                error!("Challenge server accept error: {}", e);
                                tokio::time::sleep(std::time::Duration::from_millis(100)).await;
                            }
                        }
                    }
                }
            }
        });

        self.handle = Some(handle);
        Ok(())
    }

    /// Stop the challenge server.
    pub async fn stop(&mut self) {
        self.cancel.cancel();
        if let Some(handle) = self.handle.take() {
            let _ = tokio::time::timeout(
                std::time::Duration::from_secs(5),
                handle,
            ).await;
        }
        self.challenges.clear();
        self.cancel = CancellationToken::new();
        info!("ACME challenge server stopped");
    }

    /// Handle an HTTP request for ACME challenges.
    fn handle_request(
        req: Request<Incoming>,
        challenges: &DashMap<String, String>,
    ) -> Result<Response<Full<Bytes>>, hyper::Error> {
        let path = req.uri().path();

        if let Some(token) = path.strip_prefix("/.well-known/acme-challenge/") {
            if let Some(key_auth) = challenges.get(token) {
                debug!("Serving ACME challenge for token: {}", token);
                return Ok(Response::builder()
                    .status(StatusCode::OK)
                    .header("content-type", "text/plain")
                    .body(Full::new(Bytes::from(key_auth.value().clone())))
                    .unwrap());
            }
        }

        Ok(Response::builder()
            .status(StatusCode::NOT_FOUND)
            .body(Full::new(Bytes::from("Not Found")))
            .unwrap())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_challenge_server_lifecycle() {
        let mut server = ChallengeServer::new();

        // Set a challenge before starting
        server.set_challenge("test-token".to_string(), "test-key-auth".to_string());

        // Start on a random port
        server.start(19900).await.unwrap();

        // Give server a moment to start
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;

        // Fetch the challenge
        let client = tokio::net::TcpStream::connect("127.0.0.1:19900").await.unwrap();
        let io = TokioIo::new(client);
        let (mut sender, conn) = hyper::client::conn::http1::handshake(io).await.unwrap();
        tokio::spawn(async move { let _ = conn.await; });

        let req = Request::get("/.well-known/acme-challenge/test-token")
            .body(Full::new(Bytes::new()))
            .unwrap();
        let resp = sender.send_request(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);

        // Test 404 for unknown token
        let req = Request::get("/.well-known/acme-challenge/unknown")
            .body(Full::new(Bytes::new()))
            .unwrap();
        let resp = sender.send_request(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::NOT_FOUND);

        server.stop().await;
    }
}