axum-reverse-proxy 1.3.0

A flexible and efficient reverse proxy implementation for Axum web applications
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169

use axum::{
    Router,
    body::Body,
    http::{Request, StatusCode},
    response::Json,
    routing::get,
};
use axum_reverse_proxy::ReverseProxy;
use http::header::{HeaderName, HeaderValue};
use serde_json::{Value, json};
use std::time::Duration;
use tokio::net::TcpListener;
use tower::ServiceBuilder;
use tower_http::timeout::TimeoutLayer;
use tower_http::validate_request::ValidateRequestHeaderLayer;

#[allow(deprecated)] // ValidateRequestHeaderLayer::bearer is deprecated in tower-http 0.6.8
#[tokio::test]
async fn test_proxy_with_middleware() {
    // Create a test server that echoes headers
    let app = Router::new().route(
        "/headers",
        get(|req: Request<Body>| async move {
            let headers = req
                .headers()
                .iter()
                .map(|(k, v)| (k.as_str().to_string(), v.to_str().unwrap().to_string()))
                .collect::<Vec<_>>();
            Json(json!({ "headers": headers }))
        }),
    );

    let test_listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let test_addr = test_listener.local_addr().unwrap();
    let test_server = tokio::spawn(async move {
        axum::serve(test_listener, app).await.unwrap();
    });

    // Create a reverse proxy
    let proxy = ReverseProxy::new("/", &format!("http://{test_addr}"));
    let proxy_router: Router = proxy.into();

    // Add middleware stack
    let app = proxy_router.layer(
        ServiceBuilder::new()
            // Add request timeout
            .layer(TimeoutLayer::with_status_code(
                StatusCode::REQUEST_TIMEOUT,
                Duration::from_secs(10),
            ))
            // Require API key
            .layer(ValidateRequestHeaderLayer::bearer("test-token"))
            // Add custom header
            .map_request(|mut req: Request<Body>| {
                req.headers_mut().insert(
                    HeaderName::from_static("x-custom-header"),
                    HeaderValue::from_static("custom-value"),
                );
                req
            }),
    );

    let proxy_listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let proxy_addr = proxy_listener.local_addr().unwrap();
    let proxy_server = tokio::spawn(async move {
        axum::serve(proxy_listener, app).await.unwrap();
    });

    // Create a client
    let client = reqwest::Client::new();

    // Test 1: Request without auth token should fail
    let response = client
        .get(format!("http://{proxy_addr}/headers"))
        .send()
        .await
        .unwrap();
    assert_eq!(
        response.status().as_u16(),
        StatusCode::UNAUTHORIZED.as_u16()
    );

    // Test 2: Request with auth token should succeed and include custom header
    let response = client
        .get(format!("http://{proxy_addr}/headers"))
        .header("Authorization", "Bearer test-token")
        .send()
        .await
        .unwrap();
    assert_eq!(response.status().as_u16(), StatusCode::OK.as_u16());

    let body: Value = response.json().await.unwrap();
    let headers = body.get("headers").unwrap().as_array().unwrap();

    // Verify custom header was added
    let has_custom_header = headers.iter().any(|h| {
        h.as_array()
            .unwrap()
            .first()
            .unwrap()
            .as_str()
            .unwrap()
            .eq_ignore_ascii_case("x-custom-header")
            && h.as_array()
                .unwrap()
                .get(1)
                .unwrap()
                .as_str()
                .unwrap()
                .eq_ignore_ascii_case("custom-value")
    });
    assert!(has_custom_header, "Custom header not found in response");

    // Clean up
    proxy_server.abort();
    test_server.abort();
}

#[tokio::test]
async fn test_proxy_timeout_middleware() {
    // Create a test server with artificial delay
    let app = Router::new().route(
        "/slow",
        get(|| async {
            tokio::time::sleep(Duration::from_secs(2)).await;
            "Done"
        }),
    );

    let test_listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let test_addr = test_listener.local_addr().unwrap();
    let test_server = tokio::spawn(async move {
        axum::serve(test_listener, app).await.unwrap();
    });

    // Create a reverse proxy with a short timeout
    let proxy = ReverseProxy::new("/", &format!("http://{test_addr}"));
    let proxy_router: Router = proxy.into();

    // Add timeout middleware
    let app = proxy_router.layer(TimeoutLayer::with_status_code(
        StatusCode::REQUEST_TIMEOUT,
        Duration::from_millis(100),
    ));

    let proxy_listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let proxy_addr = proxy_listener.local_addr().unwrap();
    let proxy_server = tokio::spawn(async move {
        axum::serve(proxy_listener, app).await.unwrap();
    });

    // Create a client
    let client = reqwest::Client::new();

    // Test: Request should timeout
    let response = client
        .get(format!("http://{proxy_addr}/slow"))
        .send()
        .await
        .unwrap();
    assert_eq!(
        response.status().as_u16(),
        StatusCode::REQUEST_TIMEOUT.as_u16()
    );

    // Clean up
    proxy_server.abort();
    test_server.abort();
}