doh-forwarder/tests/integration.rs

//! Integration tests for doh-forwarder.
//!
//! These tests spin up the full server on a random port and verify
//! the HTTP endpoints accept and forward DNS queries.

use axum::body::Body;
use axum::http::{header, Request, StatusCode};
use axum::routing::{get, post};
use axum::{Extension, Router};
use base64::Engine;
use hickory_proto::op::{Message, Query, ResponseCode};
use hickory_proto::rr::{Name, RecordType};
use std::sync::Arc;
use tokio::net::TcpListener;
use tower::ServiceExt;

use doh_forwarder::routes::{dns_query_get, dns_query_post, AppState};
use doh_forwarder::upstream::doh_client::DohClient;

/// Build a minimal DNS query for the given domain and record type.
fn build_dns_query(name: &str, rtype: RecordType) -> Vec<u8> {
    let mut msg = Message::new();
    msg.set_id(0x1234);
    msg.set_recursion_desired(true);
    msg.add_query(Query::query(Name::from_ascii(name).unwrap(), rtype));
    msg.to_vec().unwrap()
}

/// Create a test AppState with a dummy resolver (will fail on forward,
/// but lets us test parsing/validation/health).
fn test_state() -> AppState {
    let resolvers = vec![doh_forwarder::config::ResolverEntry {
        name: "test".into(),
        url: "http://127.0.0.1:1/dns-query".into(), // unreachable
    }];
    AppState {
        doh_client: DohClient::new(resolvers),
        rr_counter: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
    }
}

fn test_app() -> Router {
    Router::new()
        .route("/dns-query", post(dns_query_post).get(dns_query_get))
        .route("/health", get(|| async { "ok" }))
        .layer(Extension(test_state()))
}

#[tokio::test]
async fn test_health_endpoint() {
    let app = test_app();
    let req = Request::builder()
        .uri("/health")
        .body(Body::empty())
        .unwrap();

    let resp = app.oneshot(req).await.unwrap();
    assert_eq!(resp.status(), StatusCode::OK);

    let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
        .await
        .unwrap();
    assert_eq!(body.as_ref(), b"ok");
}

#[tokio::test]
async fn test_dns_query_post_valid_query() {
    let app = test_app();
    let query_bytes = build_dns_query("example.com.", RecordType::A);

    let req = Request::builder()
        .method("POST")
        .uri("/dns-query")
        .header(header::CONTENT_TYPE, "application/dns-message")
        .header(header::ACCEPT, "application/dns-message")
        .body(Body::from(query_bytes))
        .unwrap();

    let resp = app.oneshot(req).await.unwrap();
    // The upstream is unreachable, so we expect SERVFAIL (200 with DNS response)
    assert_eq!(resp.status(), StatusCode::OK);

    let content_type = resp.headers().get(header::CONTENT_TYPE).unwrap();
    assert_eq!(content_type, "application/dns-message");

    let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
        .await
        .unwrap();
    let response_msg = Message::from_vec(&body).unwrap();
    assert_eq!(response_msg.id(), 0x1234);
    assert_eq!(response_msg.response_code(), ResponseCode::ServFail);
}

#[tokio::test]
async fn test_dns_query_post_invalid_body() {
    let app = test_app();

    let req = Request::builder()
        .method("POST")
        .uri("/dns-query")
        .header(header::CONTENT_TYPE, "application/dns-message")
        .body(Body::from(vec![0x00; 3])) // too short, invalid DNS
        .unwrap();

    let resp = app.oneshot(req).await.unwrap();
    assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
}

#[tokio::test]
async fn test_dns_query_get_valid_query() {
    let app = test_app();
    let query_bytes = build_dns_query("example.com.", RecordType::A);
    let encoded = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(&query_bytes);

    let req = Request::builder()
        .uri(format!("/dns-query?dns={encoded}"))
        .body(Body::empty())
        .unwrap();

    let resp = app.oneshot(req).await.unwrap();
    assert_eq!(resp.status(), StatusCode::OK);

    let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
        .await
        .unwrap();
    let response_msg = Message::from_vec(&body).unwrap();
    assert_eq!(response_msg.id(), 0x1234);
    assert_eq!(response_msg.response_code(), ResponseCode::ServFail);
}

#[tokio::test]
async fn test_dns_query_get_invalid_base64() {
    let app = test_app();

    let req = Request::builder()
        .uri("/dns-query?dns=!!!invalid-base64!!!")
        .body(Body::empty())
        .unwrap();

    let resp = app.oneshot(req).await.unwrap();
    assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
}