obscura-server 0.2.3

A server for relaying secure messages using the Signal Protocol
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

use crate::config::Config;
use crate::core::notification::Notifier;
use crate::storage::DbPool;
use axum::{
    Router,
    body::Body,
    extract::State,
    http::{Request, StatusCode},
    middleware::{Next, from_fn_with_state},
    response::Response,
    routing::{get, post},
};
use std::sync::Arc;
use tower_governor::{GovernorLayer, governor::GovernorConfigBuilder};
use tracing::warn;

pub mod auth;
pub mod gateway;
pub mod keys;
pub mod messages;
pub mod middleware;
pub mod rate_limit;

#[derive(Clone)]
pub struct AppState {
    pub pool: DbPool,
    pub config: Config,
    pub notifier: Arc<dyn Notifier>,
    pub extractor: rate_limit::IpKeyExtractor,
}

pub fn app_router(pool: DbPool, config: Config, notifier: Arc<dyn Notifier>) -> Router {
    let extractor = rate_limit::IpKeyExtractor::new(&config.trusted_proxies);

    // Standard Tier: For general API usage
    let std_interval_ns = 1_000_000_000 / config.rate_limit_per_second.max(1);
    let standard_conf = Arc::new(
        GovernorConfigBuilder::default()
            .per_nanosecond(std_interval_ns as u64)
            .burst_size(config.rate_limit_burst)
            .key_extractor(extractor.clone())
            .finish()
            .unwrap(),
    );

    // Auth Tier: Stricter limits for expensive/sensitive registration & login
    let auth_interval_ns = 1_000_000_000 / config.auth_rate_limit_per_second.max(1);
    let auth_conf = Arc::new(
        GovernorConfigBuilder::default()
            .per_nanosecond(auth_interval_ns as u64)
            .burst_size(config.auth_rate_limit_burst)
            .key_extractor(extractor.clone())
            .finish()
            .unwrap(),
    );

    let state = AppState { pool, config, notifier, extractor };

    // Sensitive routes with strict limits
    let auth_routes = Router::new()
        .route("/accounts", post(auth::register))
        .route("/sessions", post(auth::login))
        .layer(GovernorLayer::new(auth_conf));

    // Standard routes
    let api_routes = Router::new()
        .route("/keys", post(keys::upload_keys))
        .route("/keys/{userId}", get(keys::get_pre_key_bundle))
        .route("/messages/{recipientId}", post(messages::send_message))
        .route("/gateway", get(gateway::websocket_handler))
        .layer(GovernorLayer::new(standard_conf));

    Router::new()
        .nest("/v1", auth_routes.merge(api_routes))
        .layer(from_fn_with_state(state.clone(), log_rate_limit_events))
        .with_state(state)
}

async fn log_rate_limit_events(State(state): State<AppState>, req: Request<Body>, next: Next) -> Response {
    let method = req.method().clone();

    // We must extract information BEFORE calling next.run(req), as that consumes the request.
    let path = req.uri().path().to_string();
    let headers = req.headers().clone();
    let peer_addr = req.extensions().get::<axum::extract::ConnectInfo<std::net::SocketAddr>>().map(|info| info.0.ip());

    let mut response = next.run(req).await;

    if response.status() == StatusCode::TOO_MANY_REQUESTS {
        // Use the exact same secure logic as the rate limiter to identify the client IP.
        let ip = peer_addr
            .map(|addr| state.extractor.identify_client_ip(&headers, addr).to_string())
            .unwrap_or_else(|| "unknown".into());

        warn!("Rate limit hit: client_ip={}, method={}, path={}", ip, method, path);

        // Map the internal x-ratelimit-after to the standard Retry-After header
        // for better compatibility with standard HTTP clients.
        if let Some(after) = response.headers().get("x-ratelimit-after") {
            let after = after.clone();
            response.headers_mut().insert("retry-after", after);
        }
    }

    response
}