itsi-server 0.1.1 → 0.1.18

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (184) hide show
  1. checksums.yaml +4 -4
  2. data/CHANGELOG.md +5 -0
  3. data/CODE_OF_CONDUCT.md +7 -0
  4. data/Cargo.lock +3937 -0
  5. data/Cargo.toml +7 -0
  6. data/README.md +4 -0
  7. data/Rakefile +8 -1
  8. data/_index.md +6 -0
  9. data/exe/itsi +141 -46
  10. data/ext/itsi_error/Cargo.toml +3 -0
  11. data/ext/itsi_error/src/lib.rs +98 -24
  12. data/ext/itsi_error/target/debug/build/clang-sys-da71b0344e568175/out/common.rs +355 -0
  13. data/ext/itsi_error/target/debug/build/clang-sys-da71b0344e568175/out/dynamic.rs +276 -0
  14. data/ext/itsi_error/target/debug/build/clang-sys-da71b0344e568175/out/macros.rs +49 -0
  15. data/ext/itsi_error/target/debug/build/rb-sys-49f554618693db24/out/bindings-0.9.110-mri-arm64-darwin23-3.4.2.rs +8865 -0
  16. data/ext/itsi_error/target/debug/incremental/itsi_error-1mmt5sux7jb0i/s-h510z7m8v9-0bxu7yd.lock +0 -0
  17. data/ext/itsi_error/target/debug/incremental/itsi_error-2vn3jey74oiw0/s-h5113n0e7e-1v5qzs6.lock +0 -0
  18. data/ext/itsi_error/target/debug/incremental/itsi_error-37uv9dicz7awp/s-h510ykifhe-0tbnep2.lock +0 -0
  19. data/ext/itsi_error/target/debug/incremental/itsi_error-37uv9dicz7awp/s-h510yyocpj-0tz7ug7.lock +0 -0
  20. data/ext/itsi_error/target/debug/incremental/itsi_error-37uv9dicz7awp/s-h510z0xc8g-14ol18k.lock +0 -0
  21. data/ext/itsi_error/target/debug/incremental/itsi_error-3g5qf4y7d54uj/s-h5113n0e7d-1trk8on.lock +0 -0
  22. data/ext/itsi_error/target/debug/incremental/itsi_error-3lpfftm45d3e2/s-h510z7m8r3-1pxp20o.lock +0 -0
  23. data/ext/itsi_error/target/debug/incremental/itsi_error-3o4qownhl3d7n/s-h510ykifek-1uxasnk.lock +0 -0
  24. data/ext/itsi_error/target/debug/incremental/itsi_error-3o4qownhl3d7n/s-h510yyocki-11u37qm.lock +0 -0
  25. data/ext/itsi_error/target/debug/incremental/itsi_error-3o4qownhl3d7n/s-h510z0xc93-0pmy0zm.lock +0 -0
  26. data/ext/itsi_instrument_entry/Cargo.toml +15 -0
  27. data/ext/itsi_instrument_entry/src/lib.rs +31 -0
  28. data/ext/itsi_rb_helpers/Cargo.toml +3 -0
  29. data/ext/itsi_rb_helpers/src/heap_value.rs +139 -0
  30. data/ext/itsi_rb_helpers/src/lib.rs +140 -10
  31. data/ext/itsi_rb_helpers/target/debug/build/clang-sys-da71b0344e568175/out/common.rs +355 -0
  32. data/ext/itsi_rb_helpers/target/debug/build/clang-sys-da71b0344e568175/out/dynamic.rs +276 -0
  33. data/ext/itsi_rb_helpers/target/debug/build/clang-sys-da71b0344e568175/out/macros.rs +49 -0
  34. data/ext/itsi_rb_helpers/target/debug/build/rb-sys-eb9ed4ff3a60f995/out/bindings-0.9.110-mri-arm64-darwin23-3.4.2.rs +8865 -0
  35. data/ext/itsi_rb_helpers/target/debug/incremental/itsi_rb_helpers-040pxg6yhb3g3/s-h5113n7a1b-03bwlt4.lock +0 -0
  36. data/ext/itsi_rb_helpers/target/debug/incremental/itsi_rb_helpers-131g1u4dzkt1a/s-h51113xnh3-1eik1ip.lock +0 -0
  37. data/ext/itsi_rb_helpers/target/debug/incremental/itsi_rb_helpers-131g1u4dzkt1a/s-h5111704jj-0g4rj8x.lock +0 -0
  38. data/ext/itsi_rb_helpers/target/debug/incremental/itsi_rb_helpers-1q2d3drtxrzs5/s-h5113n79yl-0bxcqc5.lock +0 -0
  39. data/ext/itsi_rb_helpers/target/debug/incremental/itsi_rb_helpers-374a9h7ovycj0/s-h51113xoox-10de2hp.lock +0 -0
  40. data/ext/itsi_rb_helpers/target/debug/incremental/itsi_rb_helpers-374a9h7ovycj0/s-h5111704w7-0vdq7gq.lock +0 -0
  41. data/ext/itsi_scheduler/Cargo.toml +24 -0
  42. data/ext/itsi_scheduler/src/itsi_scheduler/io_helpers.rs +56 -0
  43. data/ext/itsi_scheduler/src/itsi_scheduler/io_waiter.rs +44 -0
  44. data/ext/itsi_scheduler/src/itsi_scheduler/timer.rs +44 -0
  45. data/ext/itsi_scheduler/src/itsi_scheduler.rs +308 -0
  46. data/ext/itsi_scheduler/src/lib.rs +38 -0
  47. data/ext/itsi_server/Cargo.lock +2956 -0
  48. data/ext/itsi_server/Cargo.toml +72 -14
  49. data/ext/itsi_server/extconf.rb +1 -1
  50. data/ext/itsi_server/src/default_responses/html/401.html +68 -0
  51. data/ext/itsi_server/src/default_responses/html/403.html +68 -0
  52. data/ext/itsi_server/src/default_responses/html/404.html +68 -0
  53. data/ext/itsi_server/src/default_responses/html/413.html +71 -0
  54. data/ext/itsi_server/src/default_responses/html/429.html +68 -0
  55. data/ext/itsi_server/src/default_responses/html/500.html +71 -0
  56. data/ext/itsi_server/src/default_responses/html/502.html +71 -0
  57. data/ext/itsi_server/src/default_responses/html/503.html +68 -0
  58. data/ext/itsi_server/src/default_responses/html/504.html +69 -0
  59. data/ext/itsi_server/src/default_responses/html/index.html +238 -0
  60. data/ext/itsi_server/src/default_responses/json/401.json +6 -0
  61. data/ext/itsi_server/src/default_responses/json/403.json +6 -0
  62. data/ext/itsi_server/src/default_responses/json/404.json +6 -0
  63. data/ext/itsi_server/src/default_responses/json/413.json +6 -0
  64. data/ext/itsi_server/src/default_responses/json/429.json +6 -0
  65. data/ext/itsi_server/src/default_responses/json/500.json +6 -0
  66. data/ext/itsi_server/src/default_responses/json/502.json +6 -0
  67. data/ext/itsi_server/src/default_responses/json/503.json +6 -0
  68. data/ext/itsi_server/src/default_responses/json/504.json +6 -0
  69. data/ext/itsi_server/src/default_responses/mod.rs +11 -0
  70. data/ext/itsi_server/src/env.rs +43 -0
  71. data/ext/itsi_server/src/lib.rs +132 -40
  72. data/ext/itsi_server/src/prelude.rs +2 -0
  73. data/ext/itsi_server/src/ruby_types/itsi_body_proxy/big_bytes.rs +109 -0
  74. data/ext/itsi_server/src/ruby_types/itsi_body_proxy/mod.rs +143 -0
  75. data/ext/itsi_server/src/ruby_types/itsi_grpc_call.rs +344 -0
  76. data/ext/itsi_server/src/ruby_types/itsi_grpc_response_stream/mod.rs +264 -0
  77. data/ext/itsi_server/src/ruby_types/itsi_http_request.rs +345 -0
  78. data/ext/itsi_server/src/ruby_types/itsi_http_response.rs +391 -0
  79. data/ext/itsi_server/src/ruby_types/itsi_server/file_watcher.rs +225 -0
  80. data/ext/itsi_server/src/ruby_types/itsi_server/itsi_server_config.rs +375 -0
  81. data/ext/itsi_server/src/ruby_types/itsi_server.rs +83 -0
  82. data/ext/itsi_server/src/ruby_types/mod.rs +48 -0
  83. data/ext/itsi_server/src/server/binds/bind.rs +201 -0
  84. data/ext/itsi_server/src/server/binds/bind_protocol.rs +37 -0
  85. data/ext/itsi_server/src/server/binds/listener.rs +432 -0
  86. data/ext/itsi_server/src/server/binds/mod.rs +4 -0
  87. data/ext/itsi_server/src/server/binds/tls/locked_dir_cache.rs +132 -0
  88. data/ext/itsi_server/src/server/binds/tls.rs +270 -0
  89. data/ext/itsi_server/src/server/byte_frame.rs +32 -0
  90. data/ext/itsi_server/src/server/http_message_types.rs +97 -0
  91. data/ext/itsi_server/src/server/io_stream.rs +105 -0
  92. data/ext/itsi_server/src/server/lifecycle_event.rs +12 -0
  93. data/ext/itsi_server/src/server/middleware_stack/middleware.rs +165 -0
  94. data/ext/itsi_server/src/server/middleware_stack/middlewares/allow_list.rs +56 -0
  95. data/ext/itsi_server/src/server/middleware_stack/middlewares/auth_api_key.rs +87 -0
  96. data/ext/itsi_server/src/server/middleware_stack/middlewares/auth_basic.rs +86 -0
  97. data/ext/itsi_server/src/server/middleware_stack/middlewares/auth_jwt.rs +285 -0
  98. data/ext/itsi_server/src/server/middleware_stack/middlewares/cache_control.rs +142 -0
  99. data/ext/itsi_server/src/server/middleware_stack/middlewares/compression.rs +289 -0
  100. data/ext/itsi_server/src/server/middleware_stack/middlewares/cors.rs +292 -0
  101. data/ext/itsi_server/src/server/middleware_stack/middlewares/deny_list.rs +55 -0
  102. data/ext/itsi_server/src/server/middleware_stack/middlewares/error_response/default_responses.rs +190 -0
  103. data/ext/itsi_server/src/server/middleware_stack/middlewares/error_response.rs +157 -0
  104. data/ext/itsi_server/src/server/middleware_stack/middlewares/etag.rs +195 -0
  105. data/ext/itsi_server/src/server/middleware_stack/middlewares/header_interpretation.rs +82 -0
  106. data/ext/itsi_server/src/server/middleware_stack/middlewares/intrusion_protection.rs +201 -0
  107. data/ext/itsi_server/src/server/middleware_stack/middlewares/log_requests.rs +82 -0
  108. data/ext/itsi_server/src/server/middleware_stack/middlewares/max_body.rs +47 -0
  109. data/ext/itsi_server/src/server/middleware_stack/middlewares/mod.rs +87 -0
  110. data/ext/itsi_server/src/server/middleware_stack/middlewares/proxy.rs +414 -0
  111. data/ext/itsi_server/src/server/middleware_stack/middlewares/rate_limit.rs +131 -0
  112. data/ext/itsi_server/src/server/middleware_stack/middlewares/redirect.rs +76 -0
  113. data/ext/itsi_server/src/server/middleware_stack/middlewares/request_headers.rs +44 -0
  114. data/ext/itsi_server/src/server/middleware_stack/middlewares/response_headers.rs +36 -0
  115. data/ext/itsi_server/src/server/middleware_stack/middlewares/ruby_app.rs +126 -0
  116. data/ext/itsi_server/src/server/middleware_stack/middlewares/static_assets.rs +180 -0
  117. data/ext/itsi_server/src/server/middleware_stack/middlewares/static_response.rs +55 -0
  118. data/ext/itsi_server/src/server/middleware_stack/middlewares/string_rewrite.rs +163 -0
  119. data/ext/itsi_server/src/server/middleware_stack/middlewares/token_source.rs +12 -0
  120. data/ext/itsi_server/src/server/middleware_stack/mod.rs +347 -0
  121. data/ext/itsi_server/src/server/mod.rs +12 -5
  122. data/ext/itsi_server/src/server/process_worker.rs +247 -0
  123. data/ext/itsi_server/src/server/request_job.rs +11 -0
  124. data/ext/itsi_server/src/server/serve_strategy/cluster_mode.rs +342 -0
  125. data/ext/itsi_server/src/server/serve_strategy/mod.rs +30 -0
  126. data/ext/itsi_server/src/server/serve_strategy/single_mode.rs +421 -0
  127. data/ext/itsi_server/src/server/signal.rs +76 -0
  128. data/ext/itsi_server/src/server/size_limited_incoming.rs +101 -0
  129. data/ext/itsi_server/src/server/thread_worker.rs +475 -0
  130. data/ext/itsi_server/src/services/cache_store.rs +74 -0
  131. data/ext/itsi_server/src/services/itsi_http_service.rs +239 -0
  132. data/ext/itsi_server/src/services/mime_types.rs +1416 -0
  133. data/ext/itsi_server/src/services/mod.rs +6 -0
  134. data/ext/itsi_server/src/services/password_hasher.rs +83 -0
  135. data/ext/itsi_server/src/services/rate_limiter.rs +569 -0
  136. data/ext/itsi_server/src/services/static_file_server.rs +1324 -0
  137. data/ext/itsi_tracing/Cargo.toml +5 -0
  138. data/ext/itsi_tracing/src/lib.rs +315 -7
  139. data/ext/itsi_tracing/target/debug/incremental/itsi_tracing-0994n8rpvvt9m/s-h510hfz1f6-1kbycmq.lock +0 -0
  140. data/ext/itsi_tracing/target/debug/incremental/itsi_tracing-0bob7bf4yq34i/s-h5113125h5-0lh4rag.lock +0 -0
  141. data/ext/itsi_tracing/target/debug/incremental/itsi_tracing-2fcodulrxbbxo/s-h510h2infk-0hp5kjw.lock +0 -0
  142. data/ext/itsi_tracing/target/debug/incremental/itsi_tracing-2iak63r1woi1l/s-h510h2in4q-0kxfzw1.lock +0 -0
  143. data/ext/itsi_tracing/target/debug/incremental/itsi_tracing-2kk4qj9gn5dg2/s-h5113124kv-0enwon2.lock +0 -0
  144. data/ext/itsi_tracing/target/debug/incremental/itsi_tracing-2mwo0yas7dtw4/s-h510hfz1ha-1udgpei.lock +0 -0
  145. data/lib/itsi/http_request/response_status_shortcodes.rb +74 -0
  146. data/lib/itsi/http_request.rb +186 -0
  147. data/lib/itsi/http_response.rb +41 -0
  148. data/lib/itsi/passfile.rb +109 -0
  149. data/lib/itsi/server/config/dsl.rb +565 -0
  150. data/lib/itsi/server/config.rb +166 -0
  151. data/lib/itsi/server/default_app/default_app.rb +34 -0
  152. data/lib/itsi/server/default_app/index.html +115 -0
  153. data/lib/itsi/server/default_config/Itsi-rackup.rb +119 -0
  154. data/lib/itsi/server/default_config/Itsi.rb +107 -0
  155. data/lib/itsi/server/grpc/grpc_call.rb +246 -0
  156. data/lib/itsi/server/grpc/grpc_interface.rb +100 -0
  157. data/lib/itsi/server/grpc/reflection/v1/reflection_pb.rb +26 -0
  158. data/lib/itsi/server/grpc/reflection/v1/reflection_services_pb.rb +122 -0
  159. data/lib/itsi/server/rack/handler/itsi.rb +27 -0
  160. data/lib/itsi/server/rack_interface.rb +94 -0
  161. data/lib/itsi/server/route_tester.rb +107 -0
  162. data/lib/itsi/server/scheduler_interface.rb +21 -0
  163. data/lib/itsi/server/scheduler_mode.rb +10 -0
  164. data/lib/itsi/server/signal_trap.rb +29 -0
  165. data/lib/itsi/server/typed_handlers/param_parser.rb +200 -0
  166. data/lib/itsi/server/typed_handlers/source_parser.rb +55 -0
  167. data/lib/itsi/server/typed_handlers.rb +17 -0
  168. data/lib/itsi/server/version.rb +1 -1
  169. data/lib/itsi/server.rb +160 -9
  170. data/lib/itsi/standard_headers.rb +86 -0
  171. data/lib/ruby_lsp/itsi/addon.rb +111 -0
  172. data/lib/shell_completions/completions.rb +26 -0
  173. metadata +182 -25
  174. data/ext/itsi_server/src/request/itsi_request.rs +0 -143
  175. data/ext/itsi_server/src/request/mod.rs +0 -1
  176. data/ext/itsi_server/src/server/bind.rs +0 -138
  177. data/ext/itsi_server/src/server/itsi_ca/itsi_ca.crt +0 -32
  178. data/ext/itsi_server/src/server/itsi_ca/itsi_ca.key +0 -52
  179. data/ext/itsi_server/src/server/itsi_server.rs +0 -182
  180. data/ext/itsi_server/src/server/listener.rs +0 -218
  181. data/ext/itsi_server/src/server/tls.rs +0 -138
  182. data/ext/itsi_server/src/server/transfer_protocol.rs +0 -23
  183. data/ext/itsi_server/src/stream_writer/mod.rs +0 -21
  184. data/lib/itsi/request.rb +0 -39
@@ -0,0 +1,6 @@
1
+ pub mod cache_store;
2
+ pub mod itsi_http_service;
3
+ pub mod mime_types;
4
+ pub mod password_hasher;
5
+ pub mod rate_limiter;
6
+ pub mod static_file_server;
@@ -0,0 +1,83 @@
1
+ use argon2::{
2
+ password_hash::{rand_core::OsRng, PasswordHasher, PasswordVerifier, SaltString},
3
+ Argon2, PasswordHash,
4
+ };
5
+
6
+ use itsi_error::ItsiError;
7
+ use magnus::{error::Result, Value};
8
+ use serde::Deserialize;
9
+ use serde_magnus::deserialize;
10
+ use sha_crypt::{
11
+ sha256_check, sha256_simple, sha512_check, sha512_simple, Sha256Params, Sha512Params,
12
+ };
13
+
14
+ #[derive(Debug, Deserialize)]
15
+ pub enum HashAlgorithm {
16
+ #[serde(rename(deserialize = "bcrypt"))]
17
+ Bcrypt,
18
+ #[serde(rename(deserialize = "sha256"))]
19
+ Sha256Crypt,
20
+ #[serde(rename(deserialize = "sha512"))]
21
+ Sha512Crypt,
22
+ #[serde(rename(deserialize = "argon2"))]
23
+ Argon2,
24
+ #[serde(rename(deserialize = "none"))]
25
+ None,
26
+ }
27
+
28
+ pub fn create_password_hash(password: String, algo: Value) -> Result<String> {
29
+ let hash_algorithm: HashAlgorithm = deserialize(algo)?;
30
+ match hash_algorithm {
31
+ HashAlgorithm::Bcrypt => {
32
+ // Use the bcrypt crate for password hashing.
33
+ bcrypt::hash(&password, bcrypt::DEFAULT_COST)
34
+ .map_err(ItsiError::new)
35
+ .map(Ok)?
36
+ }
37
+ HashAlgorithm::Sha256Crypt => {
38
+ let params = Sha256Params::new(1000).unwrap();
39
+ let hash = sha256_simple(&password, &params)
40
+ .map_err(|_| ItsiError::new("SHA256 hashing failed"))?;
41
+ Ok(hash)
42
+ }
43
+ HashAlgorithm::Sha512Crypt => {
44
+ let params = Sha512Params::new(1000).unwrap();
45
+ let hash = sha512_simple(&password, &params)
46
+ .map_err(|_| ItsiError::new("SHA512 hashing failed"))?;
47
+ Ok(hash)
48
+ }
49
+ HashAlgorithm::Argon2 => {
50
+ let salt = SaltString::generate(&mut OsRng);
51
+ let argon2 = Argon2::default();
52
+ let password_hash = argon2
53
+ .hash_password(password.as_bytes(), &salt)
54
+ .map_err(|_| ItsiError::new("Argon2 hashing failed"))?
55
+ .to_string();
56
+ Ok(password_hash)
57
+ }
58
+ HashAlgorithm::None => Ok(format!("$none${}", password)),
59
+ }
60
+ }
61
+
62
+ pub fn verify_password_hash(password: &str, hash: &str) -> Result<bool> {
63
+ if hash.starts_with("$2a$") || hash.starts_with("$2b$") || hash.starts_with("$2y$") {
64
+ Ok(bcrypt::verify(password, hash).map_err(ItsiError::new)?)
65
+ } else if hash.starts_with("$5$") {
66
+ Ok(sha256_check(password, hash).is_ok())
67
+ } else if hash.starts_with("$6$") {
68
+ Ok(sha512_check(password, hash).is_ok())
69
+ } else if hash.starts_with("$argon2") {
70
+ let parsed_hash =
71
+ PasswordHash::new(hash).map_err(|_| ItsiError::new("Argon2 hash parsing failed"))?;
72
+ Ok(Argon2::default()
73
+ .verify_password(password.as_bytes(), &parsed_hash)
74
+ .is_ok())
75
+ } else if hash
76
+ .strip_prefix("$none$")
77
+ .is_some_and(|stripped| stripped == password)
78
+ {
79
+ Ok(true)
80
+ } else {
81
+ Err(ItsiError::new("Unsupported hash algorithm").into())
82
+ }
83
+ }
@@ -0,0 +1,569 @@
1
+ use async_trait::async_trait;
2
+ use rand::Rng;
3
+ use redis::aio::ConnectionManager;
4
+ use redis::{Client, RedisError, Script};
5
+ use serde::Deserialize;
6
+ use std::any::Any;
7
+ use std::collections::{HashMap, HashSet};
8
+ use std::sync::{Arc, LazyLock, Mutex};
9
+ use std::time::{Duration, Instant};
10
+ use tokio::sync::{Mutex as AsyncMutex, RwLock};
11
+ use tokio::time::timeout;
12
+ use url::Url;
13
+
14
+ #[derive(Debug)]
15
+ pub enum RateLimitError {
16
+ RedisError(RedisError),
17
+ RateLimitExceeded { limit: u64, count: u64, ttl: u64 },
18
+ LockError,
19
+ ConnectionTimeout,
20
+ }
21
+
22
+ impl From<RedisError> for RateLimitError {
23
+ fn from(err: RedisError) -> Self {
24
+ RateLimitError::RedisError(err)
25
+ }
26
+ }
27
+
28
+ impl std::fmt::Display for RateLimitError {
29
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
30
+ match self {
31
+ RateLimitError::RedisError(e) => write!(f, "Redis error: {}", e),
32
+ RateLimitError::RateLimitExceeded { limit, count, ttl } => {
33
+ write!(f, "Rate limit exceeded: {}/{} (ttl: {})", count, limit, ttl)
34
+ }
35
+ RateLimitError::LockError => write!(f, "Failed to acquire lock"),
36
+ RateLimitError::ConnectionTimeout => write!(f, "Connection timeout"),
37
+ }
38
+ }
39
+ }
40
+
41
+ /// A RateLimiter trait for limiting HTTP requests
42
+ #[async_trait]
43
+ pub trait RateLimiter: Send + Sync + std::fmt::Debug {
44
+ /// Increments the counter associated with `key` and sets its expiration.
45
+ /// Returns the new counter value.
46
+ ///
47
+ /// If the operation fails, returns Ok(0) to fail open.
48
+ async fn increment(&self, key: &str, timeout: Duration) -> Result<(u64, u64), RateLimitError>;
49
+
50
+ /// Checks if the rate limit is exceeded for the given key.
51
+ /// Returns Ok(current_count) if not exceeded, or Err(RateLimitExceeded) if exceeded.
52
+ ///
53
+ /// If there's an error (like connectivity issues), this will always return Ok
54
+ /// to allow the request through (fail open).
55
+ async fn check_limit(
56
+ &self,
57
+ key: &str,
58
+ limit: u64,
59
+ timeout: Duration,
60
+ ) -> Result<(u64, u64), RateLimitError>;
61
+
62
+ /// Returns self as Any for downcasting
63
+ fn as_any(&self) -> &dyn Any;
64
+ }
65
+
66
+ /// A Redis-backed rate limiter using an async connection manager.
67
+ /// This uses a TLS-enabled connection when the URL is prefixed with "rediss://".
68
+ #[derive(Clone)]
69
+ pub struct RedisRateLimiter {
70
+ connection: Arc<ConnectionManager>,
71
+ increment_script: Script,
72
+ }
73
+
74
+ impl std::fmt::Debug for RedisRateLimiter {
75
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
76
+ f.debug_struct("RedisRateLimiter").finish()
77
+ }
78
+ }
79
+
80
+ impl RedisRateLimiter {
81
+ /// Constructs a new RedisRateLimiter with a timeout.
82
+ ///
83
+ /// Use a connection URL like:
84
+ /// - Standard: "redis://host:port/db"
85
+ /// - With auth: "redis://:password@host:port/db"
86
+ /// - With TLS: "rediss://host:port/db"
87
+ /// - With TLS and auth: "rediss://:password@host:port/db"
88
+ pub async fn new(connection_url: &str) -> Result<Self, RateLimitError> {
89
+ // Set a reasonable timeout for connection attempts (5 seconds)
90
+ const CONNECTION_TIMEOUT: Duration = Duration::from_secs(5);
91
+
92
+ // Parse URL to extract auth information if provided
93
+ let url_result = Url::parse(connection_url);
94
+ if let Err(e) = url_result {
95
+ tracing::error!("Invalid Redis URL format: {}", e);
96
+ return Err(RateLimitError::RedisError(RedisError::from((
97
+ redis::ErrorKind::InvalidClientConfig,
98
+ "Invalid Redis URL format",
99
+ ))));
100
+ }
101
+ let client = Client::open(connection_url).map_err(RateLimitError::RedisError)?;
102
+ let connection_manager_result =
103
+ timeout(CONNECTION_TIMEOUT, ConnectionManager::new(client)).await;
104
+ let connection_manager = match connection_manager_result {
105
+ Ok(result) => result.map_err(RateLimitError::RedisError)?,
106
+ Err(_) => return Err(RateLimitError::ConnectionTimeout),
107
+ };
108
+
109
+ // Create the Lua script once when initializing the rate limiter
110
+ let increment_script = Script::new(
111
+ r#"
112
+ local current = redis.call('INCR', KEYS[1])
113
+ if redis.call('TTL', KEYS[1]) < 0 then
114
+ redis.call('EXPIRE', KEYS[1], ARGV[1])
115
+ end
116
+ return { current, ttl }
117
+ "#,
118
+ );
119
+
120
+ Ok(Self {
121
+ connection: Arc::new(connection_manager),
122
+ increment_script,
123
+ })
124
+ }
125
+
126
+ /// Bans an IP address for the specified duration
127
+ pub async fn ban_ip(
128
+ &self,
129
+ ip: &str,
130
+ reason: &str,
131
+ duration: Duration,
132
+ ) -> Result<(), RateLimitError> {
133
+ let ban_key = format!("ban:ip:{}", ip);
134
+ let timeout_secs = duration.as_secs();
135
+ let mut connection = (*self.connection).clone();
136
+
137
+ // Set the ban with the reason as the value
138
+ let _: () = redis::cmd("SET")
139
+ .arg(&ban_key)
140
+ .arg(reason)
141
+ .arg("EX")
142
+ .arg(timeout_secs)
143
+ .query_async(&mut connection)
144
+ .await
145
+ .map_err(RateLimitError::RedisError)?;
146
+
147
+ Ok(())
148
+ }
149
+
150
+ /// Checks if an IP address is banned
151
+ pub async fn is_banned(&self, ip: &str) -> Result<Option<String>, RateLimitError> {
152
+ let ban_key = format!("ban:ip:{}", ip);
153
+ let mut connection = (*self.connection).clone();
154
+
155
+ // Get the ban reason if it exists
156
+ let result: Option<String> = redis::cmd("GET")
157
+ .arg(&ban_key)
158
+ .query_async(&mut connection)
159
+ .await
160
+ .map_err(RateLimitError::RedisError)?;
161
+
162
+ Ok(result)
163
+ }
164
+ }
165
+
166
+ #[async_trait]
167
+ impl RateLimiter for RedisRateLimiter {
168
+ async fn increment(&self, key: &str, timeout: Duration) -> Result<(u64, u64), RateLimitError> {
169
+ let timeout_secs = timeout.as_secs();
170
+ let mut connection = (*self.connection).clone();
171
+
172
+ // Use the pre-compiled script (atomic approach)
173
+ match self
174
+ .increment_script
175
+ .key(key)
176
+ .arg(timeout_secs)
177
+ .invoke_async(&mut connection)
178
+ .await
179
+ {
180
+ Ok((count, ttl)) => Ok((count, ttl)),
181
+ Err(err) => {
182
+ // Log the error but return 0 to fail open
183
+ tracing::warn!("Redis rate limit error: {}", err);
184
+ Ok((0, timeout_secs))
185
+ }
186
+ }
187
+ }
188
+
189
+ async fn check_limit(
190
+ &self,
191
+ key: &str,
192
+ limit: u64,
193
+ timeout: Duration,
194
+ ) -> Result<(u64, u64), RateLimitError> {
195
+ match self.increment(key, timeout).await {
196
+ Ok((count, ttl)) if count <= limit => Ok((count, ttl)),
197
+ Ok((count, ttl)) if count > limit => {
198
+ Err(RateLimitError::RateLimitExceeded { limit, count, ttl })
199
+ }
200
+ // For any error or other case, fail open
201
+ _ => Ok((0, timeout.as_secs())),
202
+ }
203
+ }
204
+
205
+ fn as_any(&self) -> &dyn Any {
206
+ self
207
+ }
208
+ }
209
+
210
+ /// An entry in the in-memory rate limiter
211
+ #[derive(Debug)]
212
+ struct RateLimitEntry {
213
+ count: u64,
214
+ expires_at: Instant,
215
+ }
216
+
217
+ /// An in-memory implementation of the RateLimiter trait
218
+ #[derive(Debug)]
219
+ pub struct InMemoryRateLimiter {
220
+ entries: RwLock<HashMap<String, RateLimitEntry>>,
221
+ }
222
+
223
+ impl InMemoryRateLimiter {
224
+ /// Creates a new in-memory rate limiter
225
+ pub fn new() -> Self {
226
+ Self {
227
+ entries: RwLock::new(HashMap::new()),
228
+ }
229
+ }
230
+
231
+ /// Cleans up expired entries
232
+ async fn cleanup(&self) {
233
+ // Try to get the write lock, but fail open if we can't
234
+ if let Ok(mut entries) = self.entries.try_write() {
235
+ let now = Instant::now();
236
+ entries.retain(|_, entry| entry.expires_at > now);
237
+ }
238
+ }
239
+
240
+ /// Bans an IP address for the specified duration
241
+ pub async fn ban_ip(
242
+ &self,
243
+ ip: &str,
244
+ _: &str,
245
+ duration: Duration,
246
+ ) -> Result<(), RateLimitError> {
247
+ let now = Instant::now();
248
+ let ban_key = format!("ban:ip:{}", ip);
249
+
250
+ let mut entries = self.entries.try_write().map_err(|e| {
251
+ tracing::error!("Failed to acquire write lock: {}", e);
252
+ RateLimitError::LockError
253
+ })?;
254
+
255
+ entries.insert(
256
+ ban_key,
257
+ RateLimitEntry {
258
+ count: 1, // Use count=1 to indicate banned
259
+ expires_at: now + duration,
260
+ },
261
+ );
262
+
263
+ Ok(())
264
+ }
265
+
266
+ /// Checks if an IP address is banned
267
+ pub async fn is_banned(&self, ip: &str) -> Result<Option<String>, RateLimitError> {
268
+ let now = Instant::now();
269
+ let ban_key = format!("ban:ip:{}", ip);
270
+
271
+ let entries = self.entries.try_read().map_err(|e| {
272
+ tracing::error!("Failed to acquire read lock: {}", e);
273
+ RateLimitError::LockError
274
+ })?;
275
+
276
+ if let Some(entry) = entries.get(&ban_key) {
277
+ if entry.expires_at > now {
278
+ // IP is banned, return a generic reason since we don't store reasons
279
+ return Ok(Some("IP address banned".to_string()));
280
+ }
281
+ }
282
+
283
+ Ok(None)
284
+ }
285
+ }
286
+
287
+ impl Default for InMemoryRateLimiter {
288
+ fn default() -> Self {
289
+ Self::new()
290
+ }
291
+ }
292
+
293
+ #[async_trait]
294
+ impl RateLimiter for InMemoryRateLimiter {
295
+ async fn increment(&self, key: &str, timeout: Duration) -> Result<(u64, u64), RateLimitError> {
296
+ if rand::rng().random_bool(0.01) {
297
+ self.cleanup().await;
298
+ }
299
+
300
+ let now = Instant::now();
301
+
302
+ let mut entries = self.entries.write().await;
303
+
304
+ let entry = entries
305
+ .entry(key.to_string())
306
+ .or_insert_with(|| RateLimitEntry {
307
+ count: 0,
308
+ expires_at: now + timeout,
309
+ });
310
+
311
+ if entry.expires_at < now {
312
+ entry.expires_at = now + timeout;
313
+ entry.count = 1;
314
+ } else {
315
+ entry.count += 1;
316
+ }
317
+
318
+ let ttl = if entry.expires_at > now {
319
+ entry.expires_at.duration_since(now).as_secs()
320
+ } else {
321
+ 0
322
+ };
323
+
324
+ Ok((entry.count, ttl))
325
+ }
326
+
327
+ async fn check_limit(
328
+ &self,
329
+ key: &str,
330
+ limit: u64,
331
+ timeout: Duration,
332
+ ) -> Result<(u64, u64), RateLimitError> {
333
+ match self.increment(key, timeout).await {
334
+ Ok((count, ttl)) if count <= limit => Ok((count, ttl)),
335
+ Ok((count, ttl)) if count > limit => {
336
+ Err(RateLimitError::RateLimitExceeded { limit, count, ttl })
337
+ }
338
+ // For any error or other case, fail open
339
+ _ => Ok((0, timeout.as_secs())),
340
+ }
341
+ }
342
+
343
+ fn as_any(&self) -> &dyn Any {
344
+ self
345
+ }
346
+ }
347
+
348
+ /// Enum to represent different types of rate limiters that can ban IPs
349
+ #[derive(Debug, Clone)]
350
+ pub enum BanManager {
351
+ Redis(Arc<RedisRateLimiter>),
352
+ InMemory(Arc<InMemoryRateLimiter>),
353
+ }
354
+
355
+ impl BanManager {
356
+ /// Bans an IP address for the specified duration
357
+ pub async fn ban_ip(
358
+ &self,
359
+ ip: &str,
360
+ reason: &str,
361
+ duration: Duration,
362
+ ) -> Result<(), RateLimitError> {
363
+ match self {
364
+ BanManager::Redis(limiter) => limiter.ban_ip(ip, reason, duration).await,
365
+ BanManager::InMemory(limiter) => limiter.ban_ip(ip, reason, duration).await,
366
+ }
367
+ }
368
+
369
+ /// Checks if an IP address is banned
370
+ pub async fn is_banned(&self, ip: &str) -> Result<Option<String>, RateLimitError> {
371
+ match self {
372
+ BanManager::Redis(limiter) => limiter.is_banned(ip).await,
373
+ BanManager::InMemory(limiter) => limiter.is_banned(ip).await,
374
+ }
375
+ }
376
+ }
377
+
378
+ /// Utility function to create a rate limit key for a specific minute
379
+ pub fn create_rate_limit_key(api_key: &str, resource: &str) -> String {
380
+ // Get the current minute number (0-59)
381
+ let now = std::time::SystemTime::now()
382
+ .duration_since(std::time::UNIX_EPOCH)
383
+ .unwrap_or_default();
384
+
385
+ let minutes = now.as_secs() / 60 % 60;
386
+ format!("ratelimit:{}:{}:{}", api_key, resource, minutes)
387
+ }
388
+
389
+ /// Utility function to create a ban key for an IP address
390
+ pub fn create_ban_key(ip: &str) -> String {
391
+ format!("ban:ip:{}", ip)
392
+ }
393
+
394
+ // Global map of URL to mutex to ensure only one connection attempt per URL at a time
395
+ static CONNECTION_LOCKS: LazyLock<Mutex<HashMap<String, Arc<AsyncMutex<()>>>>> =
396
+ LazyLock::new(|| Mutex::new(HashMap::new()));
397
+
398
+ /// A global store for rate limiters, indexed by connection URL
399
+ pub struct RateLimiterStore {
400
+ redis_limiters: Mutex<HashMap<String, Arc<RedisRateLimiter>>>,
401
+ memory_limiter: Arc<InMemoryRateLimiter>,
402
+ // Track known bad Redis URLs to avoid repeated connection attempts
403
+ failed_urls: Mutex<HashSet<String>>,
404
+ }
405
+
406
+ impl RateLimiterStore {
407
+ /// Create a new store with a single in-memory rate limiter
408
+ fn new() -> Self {
409
+ Self {
410
+ redis_limiters: Mutex::new(HashMap::new()),
411
+ memory_limiter: Arc::new(InMemoryRateLimiter::new()),
412
+ failed_urls: Mutex::new(HashSet::new()),
413
+ }
414
+ }
415
+
416
+ /// Get an in-memory rate limiter
417
+ pub fn get_memory_limiter(&self) -> Arc<InMemoryRateLimiter> {
418
+ self.memory_limiter.clone()
419
+ }
420
+
421
+ /// Get a Redis rate limiter for the given connection URL, creating one if it doesn't exist
422
+ pub async fn get_redis_limiter(
423
+ &self,
424
+ connection_url: &str,
425
+ ) -> Result<Arc<RedisRateLimiter>, RateLimitError> {
426
+ // First check if this URL is known to fail
427
+ {
428
+ let failed_urls = self.failed_urls.lock().unwrap_or_else(|e| e.into_inner());
429
+ if failed_urls.contains(connection_url) {
430
+ return Err(RateLimitError::ConnectionTimeout);
431
+ }
432
+ }
433
+
434
+ // Then check if we already have a limiter for this URL
435
+ {
436
+ let limiters = self
437
+ .redis_limiters
438
+ .lock()
439
+ .unwrap_or_else(|e| e.into_inner());
440
+ if let Some(limiter) = limiters.get(connection_url) {
441
+ return Ok(limiter.clone());
442
+ }
443
+ }
444
+
445
+ // Get a dedicated mutex for this URL or create a new one if it doesn't exist
446
+ let url_mutex = {
447
+ let mut locks = CONNECTION_LOCKS.lock().unwrap_or_else(|e| e.into_inner());
448
+
449
+ // Get or create the mutex for this URL
450
+ locks
451
+ .entry(connection_url.to_string())
452
+ .or_insert_with(|| Arc::new(AsyncMutex::new(())))
453
+ .clone()
454
+ };
455
+
456
+ // Acquire the mutex with a timeout to avoid deadlocks
457
+ let lock_result = timeout(Duration::from_secs(5), url_mutex.lock()).await;
458
+ let _guard = match lock_result {
459
+ Ok(guard) => guard,
460
+ Err(_) => {
461
+ tracing::warn!("Timed out waiting for lock on URL: {}", connection_url);
462
+ return Err(RateLimitError::LockError);
463
+ }
464
+ };
465
+
466
+ // Check again if another thread created the limiter while we were waiting
467
+ {
468
+ let limiters = self
469
+ .redis_limiters
470
+ .lock()
471
+ .unwrap_or_else(|e| e.into_inner());
472
+ if let Some(limiter) = limiters.get(connection_url) {
473
+ return Ok(limiter.clone());
474
+ }
475
+ }
476
+
477
+ // Create a new limiter
478
+ tracing::info!("Initializing Redis rate limiter for {}", connection_url);
479
+ match RedisRateLimiter::new(connection_url).await {
480
+ Ok(limiter) => {
481
+ let limiter = Arc::new(limiter);
482
+
483
+ // Store it for future use
484
+ let mut limiters = self
485
+ .redis_limiters
486
+ .lock()
487
+ .unwrap_or_else(|e| e.into_inner());
488
+ limiters.insert(connection_url.to_string(), limiter.clone());
489
+
490
+ Ok(limiter)
491
+ }
492
+ Err(e) => {
493
+ tracing::error!("Failed to initialize Redis rate limiter: {}", e);
494
+ // Cache the failure
495
+ let mut failed_urls = self.failed_urls.lock().unwrap_or_else(|e| e.into_inner());
496
+ failed_urls.insert(connection_url.to_string());
497
+ Err(e)
498
+ }
499
+ }
500
+ }
501
+
502
+ /// Get a BanManager for the given RateLimiterConfig
503
+ pub async fn get_ban_manager(
504
+ &self,
505
+ config: &RateLimiterConfig,
506
+ ) -> Result<BanManager, RateLimitError> {
507
+ match config {
508
+ RateLimiterConfig::Memory => {
509
+ tracing::debug!("Using in-memory ban manager");
510
+ Ok(BanManager::InMemory(self.get_memory_limiter()))
511
+ }
512
+ RateLimiterConfig::Redis { connection_url } => {
513
+ match self.get_redis_limiter(connection_url).await {
514
+ Ok(limiter) => Ok(BanManager::Redis(limiter)),
515
+ Err(_) => Ok(BanManager::InMemory(self.get_memory_limiter())),
516
+ }
517
+ }
518
+ }
519
+ }
520
+ }
521
+
522
+ /// Global store of rate limiters
523
+ pub static RATE_LIMITER_STORE: LazyLock<RateLimiterStore> = LazyLock::new(RateLimiterStore::new);
524
+
525
+ /// Convenience function to get an in-memory rate limiter
526
+ pub fn get_memory_rate_limiter() -> Arc<impl RateLimiter> {
527
+ RATE_LIMITER_STORE.get_memory_limiter()
528
+ }
529
+
530
+ /// Convenience function to get a Redis rate limiter by connection URL
531
+ pub async fn get_redis_rate_limiter(
532
+ connection_url: &str,
533
+ ) -> Result<Arc<impl RateLimiter>, RateLimitError> {
534
+ RATE_LIMITER_STORE.get_redis_limiter(connection_url).await
535
+ }
536
+
537
+ /// Get a rate limiter based on configuration
538
+ pub async fn get_rate_limiter(
539
+ config: &RateLimiterConfig,
540
+ ) -> Result<Arc<dyn RateLimiter>, RateLimitError> {
541
+ match config {
542
+ RateLimiterConfig::Memory => Ok(get_memory_rate_limiter() as Arc<dyn RateLimiter>),
543
+ RateLimiterConfig::Redis { connection_url } => {
544
+ match get_redis_rate_limiter(connection_url).await {
545
+ Ok(limiter) => Ok(limiter as Arc<dyn RateLimiter>),
546
+ Err(_) => Ok(get_memory_rate_limiter() as Arc<dyn RateLimiter>),
547
+ }
548
+ }
549
+ }
550
+ }
551
+
552
+ /// Get a ban manager based on configuration
553
+ pub async fn get_ban_manager(config: &RateLimiterConfig) -> Result<BanManager, RateLimitError> {
554
+ RATE_LIMITER_STORE.get_ban_manager(config).await
555
+ }
556
+
557
+ /// Configuration for rate limiters
558
+ #[derive(Debug, Clone, Deserialize)]
559
+ pub enum RateLimiterConfig {
560
+ /// Use an in-memory rate limiter
561
+ #[serde(rename(deserialize = "in_memory"))]
562
+ Memory,
563
+ /// Use a Redis-backed rate limiter
564
+ #[serde(rename(deserialize = "redis"))]
565
+ Redis {
566
+ /// Connection URL, including database number if needed (e.g., "redis://localhost:6379/0")
567
+ connection_url: String,
568
+ },
569
+ }