From 14b199ad15ea7eb9231ebc296096b85d45a8d8c4 Mon Sep 17 00:00:00 2001
From: Javad Zobeidi <j.zobeidi89@gmail.com>
Date: Tue, 26 Mar 2024 13:31:15 +0330
Subject: [PATCH] Adding Throttle Middleware

---
 lib/src/exception/http_exception.dart        |  3 +-
 lib/src/exception/internal_server_error.dart |  1 -
 lib/src/exception/not_found_exception.dart   |  4 +-
 lib/src/exception/throttle_exception.dart    | 10 +++++
 lib/src/route/middleware/throttle.dart       | 39 ++++++++++++++++++++
 lib/src/route/throttle_requests.dart         | 33 +++++++++++++++++
 lib/vania.dart                               |  1 +
 7 files changed, 87 insertions(+), 4 deletions(-)
 create mode 100644 lib/src/exception/throttle_exception.dart
 create mode 100644 lib/src/route/middleware/throttle.dart
 create mode 100644 lib/src/route/throttle_requests.dart

diff --git a/lib/src/exception/http_exception.dart b/lib/src/exception/http_exception.dart
index 26ac331..9f62470 100644
--- a/lib/src/exception/http_exception.dart
+++ b/lib/src/exception/http_exception.dart
@@ -1,4 +1,5 @@
-import 'package:vania/vania.dart';
+import '../http/response/response.dart';
+import 'base_http_exception.dart';
 
 class HttpException extends BaseHttpException {
   HttpException(
diff --git a/lib/src/exception/internal_server_error.dart b/lib/src/exception/internal_server_error.dart
index 382da84..ad9db19 100644
--- a/lib/src/exception/internal_server_error.dart
+++ b/lib/src/exception/internal_server_error.dart
@@ -1,5 +1,4 @@
 import 'dart:io';
-
 import 'base_http_exception.dart';
 
 class InternalServerError extends BaseHttpException {
diff --git a/lib/src/exception/not_found_exception.dart b/lib/src/exception/not_found_exception.dart
index 5b2997d..79f8684 100644
--- a/lib/src/exception/not_found_exception.dart
+++ b/lib/src/exception/not_found_exception.dart
@@ -1,6 +1,6 @@
 import 'dart:io';
-
-import 'package:vania/vania.dart';
+import '../http/response/response.dart';
+import 'base_http_exception.dart';
 
 class NotFoundException extends BaseHttpException {
   NotFoundException({
diff --git a/lib/src/exception/throttle_exception.dart b/lib/src/exception/throttle_exception.dart
new file mode 100644
index 0000000..49ff0e6
--- /dev/null
+++ b/lib/src/exception/throttle_exception.dart
@@ -0,0 +1,10 @@
+import '../http/response/response.dart';
+import 'base_http_exception.dart';
+
+class ThrottleException extends BaseHttpException {
+  ThrottleException(
+      {required super.message,
+      required super.code,
+      super.responseType = ResponseType.json,
+      super.errorCode = 'Rate limiting'});
+}
diff --git a/lib/src/route/middleware/throttle.dart b/lib/src/route/middleware/throttle.dart
new file mode 100644
index 0000000..2b9d31b
--- /dev/null
+++ b/lib/src/route/middleware/throttle.dart
@@ -0,0 +1,39 @@
+import 'dart:io';
+
+import 'package:vania/vania.dart';
+
+import '../../exception/throttle_exception.dart';
+import '../throttle_requests.dart';
+
+class Throttle extends Middleware {
+  final int maxAttempts;
+  final Duration duration;
+
+  Throttle({
+    this.maxAttempts = 6,
+    this.duration = const Duration(seconds: 60),
+  }) {
+    throttle = ThrottleRequests(maxAttempts: maxAttempts, duration: duration);
+  }
+
+  late ThrottleRequests throttle;
+
+  @override
+  Future handle(Request req) async {
+    final clientIp = req.ip;
+    if (clientIp == null) {
+      req.response.statusCode = HttpStatus.internalServerError;
+      req.response.write('Error determining client IP');
+      await req.response.close();
+      return;
+    }
+
+    if (!throttle.request(clientIp)) {
+      throw ThrottleException(
+        message: 'Too Many Requests.',
+        code: HttpStatus.tooManyRequests,
+      );
+    }
+    next?.handle(req);
+  }
+}
diff --git a/lib/src/route/throttle_requests.dart b/lib/src/route/throttle_requests.dart
new file mode 100644
index 0000000..3c21bb5
--- /dev/null
+++ b/lib/src/route/throttle_requests.dart
@@ -0,0 +1,33 @@
+class ThrottleRequests {
+  final int maxAttempts;
+  final Duration duration;
+  final Map<String, _RateLimit> _limits = {};
+
+  ThrottleRequests({required this.maxAttempts, required this.duration});
+
+  bool request(String ip) {
+    final currentTime = DateTime.now();
+
+    _limits.putIfAbsent(ip, () => _RateLimit(0, currentTime));
+    final limit = _limits[ip]!;
+
+    if (currentTime.difference(limit.windowStart).compareTo(duration) >= 0) {
+      limit.count = 0;
+      limit.windowStart = currentTime;
+    }
+
+    if (limit.count < maxAttempts) {
+      limit.count++;
+      return true;
+    } else {
+      return false;
+    }
+  }
+}
+
+class _RateLimit {
+  int count;
+  DateTime windowStart;
+
+  _RateLimit(this.count, this.windowStart);
+}
diff --git a/lib/vania.dart b/lib/vania.dart
index 833f849..b33a1ba 100644
--- a/lib/vania.dart
+++ b/lib/vania.dart
@@ -20,6 +20,7 @@ export 'src/websocket/websocket_event.dart';
 
 export 'src/route/router.dart';
 export 'src/route/route.dart';
+export 'src/route/middleware/throttle.dart';
 
 export 'src/cache/cache_driver.dart';
 export 'src/cache/cache.dart';