systemprompt_users/repository/banned_ip/
queries.rs

1//! Banned-IP lookup and mutation queries.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use crate::error::Result;
7
8use super::BannedIpRepository;
9use super::types::{BanDuration, BanIpParams, BanIpWithMetadataParams, BannedIp};
10
11impl BannedIpRepository {
12    pub async fn is_banned(&self, ip_address: &str) -> Result<bool> {
13        let result = sqlx::query_scalar!(
14            r#"
15            SELECT EXISTS(
16                SELECT 1 FROM banned_ips
17                WHERE ip_address = $1
18                  AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP)
19            ) as "exists!"
20            "#,
21            ip_address
22        )
23        .fetch_one(&*self.pool)
24        .await?;
25
26        Ok(result)
27    }
28
29    pub async fn find_ban(&self, ip_address: &str) -> Result<Option<BannedIp>> {
30        let row = sqlx::query_as!(
31            BannedIp,
32            r#"
33            SELECT
34                ip_address,
35                reason,
36                banned_at,
37                expires_at,
38                ban_count,
39                last_offense_path,
40                last_user_agent,
41                is_permanent,
42                source_fingerprint,
43                ban_source,
44                associated_session_ids
45            FROM banned_ips
46            WHERE ip_address = $1
47              AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP)
48            "#,
49            ip_address
50        )
51        .fetch_optional(&*self.pool)
52        .await?;
53
54        Ok(row)
55    }
56
57    pub async fn ban_ip(&self, params: BanIpParams<'_>) -> Result<()> {
58        let expires_at = params.duration.to_expiry();
59        let is_permanent = matches!(params.duration, BanDuration::Permanent);
60
61        sqlx::query!(
62            r#"
63            INSERT INTO banned_ips (
64                ip_address, reason, expires_at, is_permanent,
65                source_fingerprint, ban_source
66            )
67            VALUES ($1, $2, $3, $4, $5, $6)
68            ON CONFLICT (ip_address) DO UPDATE SET
69                reason = $2,
70                expires_at = CASE
71                    WHEN banned_ips.is_permanent THEN banned_ips.expires_at
72                    ELSE COALESCE($3, banned_ips.expires_at)
73                END,
74                ban_count = banned_ips.ban_count + 1,
75                is_permanent = banned_ips.is_permanent OR $4,
76                source_fingerprint = COALESCE($5, banned_ips.source_fingerprint),
77                ban_source = $6
78            "#,
79            params.ip_address,
80            params.reason,
81            expires_at,
82            is_permanent,
83            params.source_fingerprint,
84            params.ban_source
85        )
86        .execute(&*self.write_pool)
87        .await?;
88
89        Ok(())
90    }
91
92    pub async fn ban_ip_with_metadata(&self, params: BanIpWithMetadataParams<'_>) -> Result<()> {
93        let expires_at = params.duration.to_expiry();
94        let is_permanent = matches!(params.duration, BanDuration::Permanent);
95        let session_ids: Option<Vec<String>> = params.session_id.map(|s| vec![s.to_owned()]);
96
97        sqlx::query!(
98            r#"
99            INSERT INTO banned_ips (
100                ip_address, reason, expires_at, is_permanent,
101                source_fingerprint, ban_source, last_offense_path,
102                last_user_agent, associated_session_ids
103            )
104            VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
105            ON CONFLICT (ip_address) DO UPDATE SET
106                reason = $2,
107                expires_at = CASE
108                    WHEN banned_ips.is_permanent THEN banned_ips.expires_at
109                    ELSE COALESCE($3, banned_ips.expires_at)
110                END,
111                ban_count = banned_ips.ban_count + 1,
112                is_permanent = banned_ips.is_permanent OR $4,
113                source_fingerprint = COALESCE($5, banned_ips.source_fingerprint),
114                ban_source = $6,
115                last_offense_path = COALESCE($7, banned_ips.last_offense_path),
116                last_user_agent = COALESCE($8, banned_ips.last_user_agent),
117                associated_session_ids = CASE
118                    WHEN $9::TEXT[] IS NOT NULL
119                    THEN array_cat(COALESCE(banned_ips.associated_session_ids, '{}'::TEXT[]), $9)
120                    ELSE banned_ips.associated_session_ids
121                END
122            "#,
123            params.ip_address,
124            params.reason,
125            expires_at,
126            is_permanent,
127            params.source_fingerprint,
128            params.ban_source,
129            params.offense_path,
130            params.user_agent,
131            session_ids.as_deref()
132        )
133        .execute(&*self.write_pool)
134        .await?;
135
136        Ok(())
137    }
138
139    pub async fn unban_ip(&self, ip_address: &str) -> Result<bool> {
140        let result = sqlx::query!(
141            r#"
142            DELETE FROM banned_ips
143            WHERE ip_address = $1
144            "#,
145            ip_address
146        )
147        .execute(&*self.write_pool)
148        .await?;
149
150        Ok(result.rows_affected() > 0)
151    }
152
153    pub async fn cleanup_expired(&self) -> Result<u64> {
154        let result = sqlx::query!(
155            r#"
156            DELETE FROM banned_ips
157            WHERE expires_at IS NOT NULL
158              AND expires_at < CURRENT_TIMESTAMP
159              AND NOT is_permanent
160            "#
161        )
162        .execute(&*self.write_pool)
163        .await?;
164
165        Ok(result.rows_affected())
166    }
167}
systemprompt_users/repository/banned_ip/queries.rs

systemprompt_users/repository/banned_ip/
queries.rs
