this repo has no description
lewis.moe
1use crate::appview::DidResolver;
2use crate::cache::{Cache, DistributedRateLimiter, create_cache};
3use crate::circuit_breaker::CircuitBreakers;
4use crate::config::AuthConfig;
5use crate::rate_limit::RateLimiters;
6use crate::repo::PostgresBlockStore;
7use crate::storage::{BlobStorage, S3BlobStorage};
8use crate::sync::firehose::SequencedEvent;
9use sqlx::PgPool;
10use std::sync::Arc;
11use tokio::sync::broadcast;
12
13#[derive(Clone)]
14pub struct AppState {
15 pub db: PgPool,
16 pub block_store: PostgresBlockStore,
17 pub blob_store: Arc<dyn BlobStorage>,
18 pub firehose_tx: broadcast::Sender<SequencedEvent>,
19 pub rate_limiters: Arc<RateLimiters>,
20 pub circuit_breakers: Arc<CircuitBreakers>,
21 pub cache: Arc<dyn Cache>,
22 pub distributed_rate_limiter: Arc<dyn DistributedRateLimiter>,
23 pub did_resolver: Arc<DidResolver>,
24}
25
26pub enum RateLimitKind {
27 Login,
28 AccountCreation,
29 PasswordReset,
30 ResetPassword,
31 RefreshSession,
32 OAuthToken,
33 OAuthAuthorize,
34 OAuthPar,
35 OAuthIntrospect,
36 AppPassword,
37 EmailUpdate,
38}
39
40impl RateLimitKind {
41 fn key_prefix(&self) -> &'static str {
42 match self {
43 Self::Login => "login",
44 Self::AccountCreation => "account_creation",
45 Self::PasswordReset => "password_reset",
46 Self::ResetPassword => "reset_password",
47 Self::RefreshSession => "refresh_session",
48 Self::OAuthToken => "oauth_token",
49 Self::OAuthAuthorize => "oauth_authorize",
50 Self::OAuthPar => "oauth_par",
51 Self::OAuthIntrospect => "oauth_introspect",
52 Self::AppPassword => "app_password",
53 Self::EmailUpdate => "email_update",
54 }
55 }
56
57 fn limit_and_window_ms(&self) -> (u32, u64) {
58 match self {
59 Self::Login => (10, 60_000),
60 Self::AccountCreation => (10, 3_600_000),
61 Self::PasswordReset => (5, 3_600_000),
62 Self::ResetPassword => (10, 60_000),
63 Self::RefreshSession => (60, 60_000),
64 Self::OAuthToken => (30, 60_000),
65 Self::OAuthAuthorize => (10, 60_000),
66 Self::OAuthPar => (30, 60_000),
67 Self::OAuthIntrospect => (30, 60_000),
68 Self::AppPassword => (10, 60_000),
69 Self::EmailUpdate => (5, 3_600_000),
70 }
71 }
72}
73
74impl AppState {
75 pub async fn new(db: PgPool) -> Self {
76 AuthConfig::init();
77
78 let block_store = PostgresBlockStore::new(db.clone());
79 let blob_store = S3BlobStorage::new().await;
80
81 let firehose_buffer_size: usize = std::env::var("FIREHOSE_BUFFER_SIZE")
82 .ok()
83 .and_then(|v| v.parse().ok())
84 .unwrap_or(10000);
85
86 let (firehose_tx, _) = broadcast::channel(firehose_buffer_size);
87 let rate_limiters = Arc::new(RateLimiters::new());
88 let circuit_breakers = Arc::new(CircuitBreakers::new());
89 let (cache, distributed_rate_limiter) = create_cache().await;
90 let did_resolver = Arc::new(DidResolver::new());
91
92 Self {
93 db,
94 block_store,
95 blob_store: Arc::new(blob_store),
96 firehose_tx,
97 rate_limiters,
98 circuit_breakers,
99 cache,
100 distributed_rate_limiter,
101 did_resolver,
102 }
103 }
104
105 pub fn with_rate_limiters(mut self, rate_limiters: RateLimiters) -> Self {
106 self.rate_limiters = Arc::new(rate_limiters);
107 self
108 }
109
110 pub fn with_circuit_breakers(mut self, circuit_breakers: CircuitBreakers) -> Self {
111 self.circuit_breakers = Arc::new(circuit_breakers);
112 self
113 }
114
115 pub async fn check_rate_limit(&self, kind: RateLimitKind, client_ip: &str) -> bool {
116 if std::env::var("DISABLE_RATE_LIMITING").is_ok() {
117 return true;
118 }
119
120 let key = format!("{}:{}", kind.key_prefix(), client_ip);
121 let limiter_name = kind.key_prefix();
122 let (limit, window_ms) = kind.limit_and_window_ms();
123
124 if !self
125 .distributed_rate_limiter
126 .check_rate_limit(&key, limit, window_ms)
127 .await
128 {
129 crate::metrics::record_rate_limit_rejection(limiter_name);
130 return false;
131 }
132
133 let limiter = match kind {
134 RateLimitKind::Login => &self.rate_limiters.login,
135 RateLimitKind::AccountCreation => &self.rate_limiters.account_creation,
136 RateLimitKind::PasswordReset => &self.rate_limiters.password_reset,
137 RateLimitKind::ResetPassword => &self.rate_limiters.reset_password,
138 RateLimitKind::RefreshSession => &self.rate_limiters.refresh_session,
139 RateLimitKind::OAuthToken => &self.rate_limiters.oauth_token,
140 RateLimitKind::OAuthAuthorize => &self.rate_limiters.oauth_authorize,
141 RateLimitKind::OAuthPar => &self.rate_limiters.oauth_par,
142 RateLimitKind::OAuthIntrospect => &self.rate_limiters.oauth_introspect,
143 RateLimitKind::AppPassword => &self.rate_limiters.app_password,
144 RateLimitKind::EmailUpdate => &self.rate_limiters.email_update,
145 };
146
147 let ok = limiter.check_key(&client_ip.to_string()).is_ok();
148 if !ok {
149 crate::metrics::record_rate_limit_rejection(limiter_name);
150 }
151 ok
152 }
153}