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

use axum::{
    Router,
    body::Body,
    extract::Json,
    http::{Request, StatusCode},
    routing::get,
};
use axum_reverse_proxy::BalancedProxy;
use serde_json::Value;
use serde_json::json;
use tokio::net::TcpListener;

#[tokio::test]
async fn test_round_robin_distribution() {
    let app1 = Router::new().route("/", get(|| async { Json(json!({"server": 1})) }));
    let listener1 = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr1 = listener1.local_addr().unwrap();
    let server1 = tokio::spawn(async move { axum::serve(listener1, app1).await.unwrap() });

    let app2 = Router::new().route("/", get(|| async { Json(json!({"server": 2})) }));
    let listener2 = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr2 = listener2.local_addr().unwrap();
    let server2 = tokio::spawn(async move { axum::serve(listener2, app2).await.unwrap() });
    let proxy = BalancedProxy::new(
        String::from("/"),
        vec![format!("http://{}", addr1), format!("http://{}", addr2)],
    );
    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_server = tokio::spawn(async move { axum::serve(proxy_listener, app).await.unwrap() });

    let client = reqwest::Client::new();
    for expected in [1, 2, 1, 2] {
        let res = client
            .get(format!("http://{proxy_addr}/"))
            .send()
            .await
            .unwrap();
        assert_eq!(res.status().as_u16(), StatusCode::OK.as_u16());
        let body: Value = res.json().await.unwrap();
        assert_eq!(body["server"], expected);
    }

    proxy_server.abort();
    server1.abort();
    server2.abort();
}

#[tokio::test]
async fn test_balanced_proxy_no_upstreams_returns_503() {
    use axum::body::Body;
    use tower::Service;

    let mut proxy = BalancedProxy::new(String::from("/"), Vec::<String>::new());

    let req = axum::http::Request::builder()
        .method("GET")
        .uri("/test")
        .body(Body::empty())
        .unwrap();

    let response = proxy.call(req).await.unwrap();
    assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);

    let body_bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
        .await
        .unwrap();
    let body_str = String::from_utf8(body_bytes.to_vec()).unwrap();
    assert!(body_str.contains("No upstream services available"));
}

#[tokio::test]
async fn test_balanced_proxy_path_stripping() {
    // Echo server that returns the path it received
    let echo = get(|req: Request<Body>| async move {
        let path = req.uri().path();
        Json(json!({ "received_path": path }))
    });

    // First upstream
    let app1 = Router::new()
        .route("/", echo.clone())
        .route("/{*path}", echo.clone());
    let listener1 = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr1 = listener1.local_addr().unwrap();
    let server1 = tokio::spawn(async move { axum::serve(listener1, app1).await.unwrap() });

    // Second upstream
    let app2 = Router::new()
        .route("/", echo.clone())
        .route("/{*path}", echo);
    let listener2 = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr2 = listener2.local_addr().unwrap();
    let server2 = tokio::spawn(async move { axum::serve(listener2, app2).await.unwrap() });

    // Balanced proxy that mounts on /api
    let proxy = BalancedProxy::new(
        String::from("/api"),
        vec![format!("http://{}", addr1), format!("http://{}", addr2)],
    );
    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_server = tokio::spawn(async move { axum::serve(proxy_listener, app).await.unwrap() });

    let client = reqwest::Client::new();
    let test_paths = ["/foo", "/bar/baz", "/baz"];

    for path in test_paths {
        let url = format!("http://{proxy_addr}/api{path}");

        let res = client.get(&url).send().await.unwrap();
        assert_eq!(res.status().as_u16(), StatusCode::OK.as_u16());
        let body: Value = res.json().await.unwrap();
        assert_eq!(body["received_path"], path);
    }

    proxy_server.abort();
    server1.abort();
    server2.abort();
}