revolt_database/util/
shield.rs

1use reqwest::Client;
2use revolt_config::config;
3use revolt_result::Result;
4use serde::{Deserialize, Serialize};
5use std::{collections::HashMap, sync::LazyLock};
6use crate::util::ip;
7
8static CLIENT: LazyLock<Client> = LazyLock::new(Client::new);
9
10#[derive(Serialize, Deserialize, Default, Debug)]
11pub struct ShieldValidationInput {
12    /// Remote user IP
13    pub ip: Option<String>,
14
15    /// User provided email
16    pub email: Option<String>,
17
18    /// Request headers
19    pub headers: Option<HashMap<String, String>>,
20
21    /// Skip alerts and monitoring for this request
22    pub dry_run: bool,
23}
24
25#[derive(Serialize, Deserialize)]
26pub struct ValidationResult {
27    /// Whether this request was blocked
28    blocked: bool,
29
30    /// Reasons for the request being blocked
31    reasons: Vec<String>,
32}
33
34pub async fn validate_shield(input: ShieldValidationInput) -> Result<()> {
35    let shield = config().await.api.security.shield;
36
37    if !shield.host.is_empty() {
38        if let Ok(response) = CLIENT
39            .post(format!("{}/validate", &shield.host))
40            .json(&input)
41            .header("Authorization", &shield.key)
42            .send()
43            .await
44        {
45            let result = response
46                .json::<ValidationResult>()
47                .await
48                .map_err(|_| create_error!(InternalError))?;
49
50            if result.blocked {
51                return Err(create_error!(BlockedByShield));
52            }
53        }
54    }
55
56    Ok(())
57}
58
59#[cfg(feature = "rocket-impl")]
60#[async_trait]
61impl<'r> rocket::request::FromRequest<'r> for ShieldValidationInput {
62    type Error = revolt_result::Error;
63
64    #[allow(clippy::collapsible_match)]
65    async fn from_request(
66        request: &'r rocket::Request<'_>,
67    ) -> rocket::request::Outcome<Self, Self::Error> {
68        rocket::request::Outcome::Success(ShieldValidationInput {
69            ip: Some(ip::rocket::to_real_ip(request).await),
70            headers: Some(
71                request
72                    .headers()
73                    .iter()
74                    .map(|entry| (entry.name.to_string(), entry.value.to_string()))
75                    .collect(),
76            ),
77            ..Default::default()
78        })
79    }
80}
81
82#[cfg(feature = "rocket-impl")]
83impl<'r> revolt_rocket_okapi::request::OpenApiFromRequest<'r> for ShieldValidationInput {
84    fn from_request_input(
85        _gen: &mut revolt_rocket_okapi::r#gen::OpenApiGenerator,
86        _name: String,
87        _required: bool,
88    ) -> revolt_rocket_okapi::Result<revolt_rocket_okapi::request::RequestHeaderInput> {
89        Ok(revolt_rocket_okapi::request::RequestHeaderInput::None)
90    }
91}
92#[cfg(feature = "axum-impl")]
93#[async_trait]
94impl<S> axum::extract::FromRequestParts<S> for ShieldValidationInput {
95    type Rejection = axum::Json<revolt_result::Error> ;
96
97    async fn from_request_parts(
98        parts: &mut axum::http::request::Parts,
99        _state: &S,
100    ) -> Result<Self, Self::Rejection> {
101        Ok(ShieldValidationInput {
102            ip: Some(ip::axum::to_real_ip(parts).await),
103            headers: Some(
104                parts
105                    .headers
106                    .iter()
107                    .map(|(name, value)| {
108                        (
109                            name.to_string(),
110                            value.to_str().map(|s| s.to_string()).unwrap_or_default(),
111                        )
112                    })
113                    .collect(),
114            ),
115            ..Default::default()
116        })
117    }
118}
revolt_database/util/shield.rs

revolt_database/util/
shield.rs
