fix: store 'block_until' in persistent storage (#133)

ChainSafe · Feb 3, 2025 · b49448a · b49448a
1 parent 3dd52e9
commit b49448a
Showing 1 changed file with 16 additions and 4 deletions.
diff --git a/src/rate_limiter.rs b/src/rate_limiter.rs
@@ -4,6 +4,7 @@ use worker::*;
 #[durable_object]
 pub struct RateLimiter {
     state: State,
+    #[allow(unused)]
     block_until: DateTime<Utc>,
 }
 
@@ -18,21 +19,32 @@ impl DurableObject for RateLimiter {
 
     async fn fetch(&mut self, _req: Request) -> Result<Response> {
         let now = Utc::now();
+        let block_until = self
+            .state
+            .storage()
+            .get("block_until")
+            .await
+            .map(|v| DateTime::<Utc>::from_timestamp(v, 0).unwrap_or_default())
+            .unwrap_or(Utc::now());
         console_log!(
             "Rate limiter invoked: now={:?}, block_until={:?}, may_sign={:?}",
             now,
-            self.block_until,
-            self.block_until <= now
+            block_until,
+            block_until <= now
         );
-        if self.block_until <= now {
+        if block_until <= now {
             // This Durable Object will be deleted after the alarm is triggered
             self.state
                 .storage()
                 .set_alarm(std::time::Duration::from_secs(
                     crate::constants::RATE_LIMIT_SECONDS as u64 + 1,
                 ))
                 .await?;
-            self.block_until = now + Duration::seconds(crate::constants::RATE_LIMIT_SECONDS);
+            let block_until = now + Duration::seconds(crate::constants::RATE_LIMIT_SECONDS);
+            self.state
+                .storage()
+                .put("block_until", block_until.timestamp())
+                .await?;
 
             Response::from_json(&true)
         } else {