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
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210

use axum::{Router, body::Body, http::Request, response::Json, routing::get};
use axum_reverse_proxy::ReverseProxy;
use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
use serde_json::{Value, json};
use std::{collections::HashMap, sync::Arc, time::Duration};
use tokio::{net::TcpListener, sync::Notify};

async fn echo_headers(req: Request<Body>) -> Json<Value> {
    let headers = req.headers().clone();
    Json(json!({ "headers": headers.iter().map(|(k, v)| {
        (k.as_str().to_string(), v.to_str().unwrap().to_string())
    }).collect::<HashMap<String, String>>() }))
}

#[tokio::test]
async fn test_proxy_header_handling() {
    // Create a test server that echoes headers
    let app = Router::new().route("/headers", get(echo_headers));

    let test_listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let test_addr = test_listener.local_addr().unwrap();
    let server_ready = Arc::new(Notify::new());
    let server_ready_clone = server_ready.clone();

    let test_server = tokio::spawn(async move {
        server_ready_clone.notify_one();
        axum::serve(test_listener, app).await.unwrap();
    });

    // Wait for test server to be ready
    server_ready.notified().await;

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

    let proxy_listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let proxy_addr = proxy_listener.local_addr().unwrap();
    let proxy_ready = Arc::new(Notify::new());
    let proxy_ready_clone = proxy_ready.clone();

    let proxy_server = tokio::spawn(async move {
        proxy_ready_clone.notify_one();
        axum::serve(proxy_listener, app).await.unwrap();
    });

    // Wait for proxy server to be ready
    proxy_ready.notified().await;

    // Create a client with a reasonable timeout
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(5))
        .build()
        .unwrap();

    // Test with a reasonable number of headers (10 instead of 50)
    let mut headers = HeaderMap::new();
    for i in 0..10 {
        headers.insert(
            HeaderName::from_bytes(format!("x-test-{i}").as_bytes()).unwrap(),
            HeaderValue::from_str(&format!("value-{i}")).unwrap(),
        );
    }

    // Use tokio::time::timeout for the test operations
    let test_result = tokio::time::timeout(Duration::from_secs(10), async {
        println!("Sending request to proxy...");
        let response = match client
            .get(format!("http://{proxy_addr}/headers"))
            .headers(headers.clone())
            .send()
            .await
        {
            Ok(r) => r,
            Err(e) => {
                println!("Failed to send request: {e}");
                return Err(e);
            }
        };

        println!("Got response with status: {}", response.status());
        assert_eq!(response.status().as_u16(), 200);

        let body = match response.json::<Value>().await {
            Ok(b) => {
                println!("Response body: {b}");
                b
            }
            Err(e) => {
                println!("Failed to parse response as JSON: {e}");
                return Err(e);
            }
        };

        let response_headers = body["headers"].as_object().unwrap();

        // Verify headers were forwarded (excluding host header)
        for (key, value) in headers.iter() {
            if key != "host" {
                println!("Checking header {}: {}", key, value.to_str().unwrap());
                assert_eq!(
                    response_headers[key.as_str()].as_str().unwrap(),
                    value.to_str().unwrap()
                );
            }
        }

        Ok(())
    })
    .await;

    // Clean up servers first
    proxy_server.abort();
    test_server.abort();

    // Then check the test result
    match test_result {
        Ok(result) => {
            if let Err(e) = result {
                panic!("Test failed with error: {e}");
            }
        }
        Err(_) => panic!("Test timed out after 10 seconds"),
    }
}

#[tokio::test]
async fn test_proxy_special_headers() {
    // Create a test server that echoes headers
    let app = Router::new().route("/headers", get(echo_headers));

    let test_listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let test_addr = test_listener.local_addr().unwrap();
    let server_ready = Arc::new(Notify::new());
    let server_ready_clone = server_ready.clone();

    let test_server = tokio::spawn(async move {
        server_ready_clone.notify_one();
        axum::serve(test_listener, app).await.unwrap();
    });

    // Wait for test server to be ready
    server_ready.notified().await;

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

    let proxy_listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let proxy_addr = proxy_listener.local_addr().unwrap();
    let proxy_ready = Arc::new(Notify::new());
    let proxy_ready_clone = proxy_ready.clone();

    let proxy_server = tokio::spawn(async move {
        proxy_ready_clone.notify_one();
        axum::serve(proxy_listener, app).await.unwrap();
    });

    // Wait for proxy server to be ready
    proxy_ready.notified().await;

    // Create a client with a reasonable timeout
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(5))
        .build()
        .unwrap();

    // Use tokio::time::timeout for the test operations
    let test_result = tokio::time::timeout(Duration::from_secs(10), async {
        println!("Testing special headers...");
        let response = client
            .get(format!("http://{proxy_addr}/headers"))
            .header("X-Forwarded-For", "192.168.1.1")
            .header("X-Real-IP", "192.168.1.1")
            .header("X-Request-ID", "test-request-1")
            .header("Accept-Encoding", "gzip, deflate")
            .send()
            .await?;

        println!("Special headers response status: {}", response.status());
        assert_eq!(response.status().as_u16(), 200);

        let body = response.json::<Value>().await?;
        println!("Response body: {body}");

        // Verify the headers were forwarded
        let headers = body["headers"].as_object().unwrap();
        assert_eq!(headers["x-forwarded-for"], "192.168.1.1");
        assert_eq!(headers["x-real-ip"], "192.168.1.1");
        assert_eq!(headers["x-request-id"], "test-request-1");
        assert!(headers.contains_key("accept-encoding"));

        Ok::<(), reqwest::Error>(())
    })
    .await;

    // Clean up servers first
    proxy_server.abort();
    test_server.abort();

    // Then check the test result
    match test_result {
        Ok(result) => {
            if let Err(e) = result {
                panic!("Test failed with error: {e}");
            }
        }
        Err(_) => panic!("Test timed out after 10 seconds"),
    }
}