provider-sample-url-finder/url_finder/src/routes.rs at 39718b64ff9ea1cc28b1b1341833aaf2afbbfb17 · fidlabs/provider-sample-url-finder · GitHub

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
use std::{sync::Arc, time::Duration};

use axum::{
    body::Body,
    http::Response,
    middleware,
    response::IntoResponse,
    routing::{get, post},
    Router,
};
use common::api_response::*;
use tower_governor::{
    governor::GovernorConfigBuilder, key_extractor::SmartIpKeyExtractor, GovernorError,
    GovernorLayer,
};
use utoipa::OpenApi;
use utoipa_swagger_ui::SwaggerUi;

use crate::{api::*, AppState};

fn too_many_requests_error_handler(error: GovernorError) -> Response<Body> {
    tracing::error!("Rate limit error: {:?}", error);

    match error {
        GovernorError::TooManyRequests { .. } => {
            too_many_requests("Rate limit exceeded").into_response()
        }
        _ => internal_server_error("Rate limit error").into_response(),
    }
}

pub fn create_routes(app_state: Arc<AppState>) -> Router<Arc<AppState>> {
    // more strict rate limiting for the sync routes
    let governor_secure = Arc::new(
        GovernorConfigBuilder::default()
            .per_millisecond(30)
            .burst_size(30)
            .key_extractor(SmartIpKeyExtractor)
            .error_handler(too_many_requests_error_handler)
            .finish()
            .unwrap(),
    );

    // less strict rate limiting for the async routes that will have internal queue anyway
    let governor_async = Arc::new(
        GovernorConfigBuilder::default()
            .per_millisecond(30)
            .burst_size(30)
            .key_extractor(SmartIpKeyExtractor)
            .error_handler(too_many_requests_error_handler)
            .finish()
            .unwrap(),
    );

    let governor_secure_limiter = governor_secure.limiter().clone();
    let governor_async_limiter = governor_async.limiter().clone();

    let interval = Duration::from_secs(60);

    // background task to clean up
    tokio::spawn(async move {
        loop {
            tokio::time::sleep(interval).await;

            tracing::debug!(
                "rate governor_secure_limiter storage size: {}",
                governor_secure_limiter.len()
            );
            tracing::debug!(
                "rate governor_async_limiter storage size: {}",
                governor_async_limiter.len()
            );

            governor_secure_limiter.retain_recent();
            governor_async_limiter.retain_recent();
        }
    });

    let swagger_routes = SwaggerUi::new("/").url("/api-doc/openapi.json", ApiDoc::openapi());

    let sync_routes = Router::new()
        .route("/url/find/:provider", get(handle_find_url_sp))
        .route(
            "/url/find/:provider/:client",
            get(handle_find_url_sp_client),
        )
        .route(
            "/url/retrievability/:provider/:client",
            get(handle_find_retri_by_client_and_sp),
        )
        .route(
            "/url/retrievability/:provider",
            get(handle_find_retri_by_sp),
        )
        .layer(middleware::from_fn_with_state(
            app_state.clone(),
            cache_middleware,
        ))
        .layer(GovernorLayer {
            config: governor_secure.clone(),
        });

    let async_routes = Router::new()
        .route("/jobs/:id", get(handle_get_job))
        .route("/jobs", post(handle_create_job))
        .layer(GovernorLayer {
            config: governor_async.clone(),
        });

    let healthcheck_route = Router::new()
        .route("/healthcheck", get(handle_healthcheck))
        .layer(GovernorLayer {
            config: governor_secure.clone(),
        });

    Router::new()
        .merge(swagger_routes)
        .merge(sync_routes)
        .merge(async_routes)
        .merge(healthcheck_route.clone())
}