feat: Improve CORS and rate limit middlewares

joanrodas · Oct 6, 2024 · a9007be · a9007be
1 parent d163325
commit a9007be
Show file tree

Hide file tree

Showing 2 changed files with 76 additions and 22 deletions.
diff --git a/PluboRoutes/Middleware/CorsMiddleware.php b/PluboRoutes/Middleware/CorsMiddleware.php
@@ -7,31 +7,43 @@ class CorsMiddleware implements MiddlewareInterface
     private $allowedOrigins;
     private $allowedMethods;
     private $allowedHeaders;
+    private $maxAge = 3600;
 
     public function __construct($allowedOrigins = '*', $allowedMethods = ['GET', 'POST', 'OPTIONS'], $allowedHeaders = ['Content-Type', 'Authorization'])
     {
-        $this->allowedOrigins = $allowedOrigins;
-        $this->allowedMethods = $allowedMethods;
-        $this->allowedHeaders = $allowedHeaders;
+        $this->allowedOrigins = is_array($allowedOrigins) ? $allowedOrigins : [$allowedOrigins];
+        $this->allowedMethods = array_map('strtoupper', $allowedMethods);
+        $this->allowedHeaders = array_map('strtolower', $allowedHeaders);
     }
 
     public function handle($request, $next)
     {
-        // Set CORS headers
-        header("Access-Control-Allow-Origin: {$this->allowedOrigins}");
-        header("Access-Control-Allow-Methods: " . implode(', ', $this->allowedMethods));
-        header("Access-Control-Allow-Headers: " . implode(', ', $this->allowedHeaders));
-
-        // If it's a preflight request (OPTIONS), allow it without running further middleware
-        if ($_SERVER['REQUEST_METHOD'] === 'OPTIONS') {
-            exit;
+        $origin = $request->get_header('origin');
+
+        if ($this->isAllowedOrigin($origin)) {
+            header("Access-Control-Allow-Origin: {$origin}");
+            header('Access-Control-Allow-Credentials: true');
+            header("Access-Control-Allow-Methods: " . implode(', ', $this->allowedMethods));
+            header("Access-Control-Allow-Headers: " . implode(', ', $this->allowedHeaders));
+            header("Access-Control-Max-Age: {$this->maxAge}");
+        }
+
+        if ($request->get_method() === 'OPTIONS') {
+            return new \WP_REST_Response(null, 204);
         }
 
-        // Enforce allowed methods for the actual request
-        if (!in_array($_SERVER['REQUEST_METHOD'], $this->allowedMethods)) {
-            return new \WP_REST_Response(['error' => 'Method not allowed'], 405); // 405 Method Not Allowed
+        if (!in_array($request->get_method(), $this->allowedMethods)) {
+            return new \WP_REST_Response(['error' => 'Method not allowed'], 405);
         }
 
         return $next($request);
     }
+
+    protected function isAllowedOrigin($origin)
+    {
+        if (in_array('*', $this->allowedOrigins)) {
+            return true;
+        }
+        return in_array($origin, $this->allowedOrigins);
+    }
 }
diff --git a/PluboRoutes/Middleware/RateLimitMiddleware.php b/PluboRoutes/Middleware/RateLimitMiddleware.php
@@ -6,27 +6,69 @@ class RateLimitMiddleware implements MiddlewareInterface
 {
     private $maxRequests;
     private $windowTime;
+    private $prefix = 'rate_limit_';
+    private $type;
 
-    public function __construct($maxRequests = 10, $windowTime = 60)
+    public function __construct($maxRequests = 10, $windowTime = 60, $type = 'ip')
     {
-        $this->maxRequests = $maxRequests;
-        $this->windowTime = $windowTime;
+        $this->maxRequests = max(1, intval($maxRequests));
+        $this->windowTime = max(1, intval($windowTime));
+        $this->type = in_array($type, ['ip', 'user', 'endpoint']) ? $type : 'ip';
     }
 
     public function handle($request, $next)
     {
-        $ip = $_SERVER['REMOTE_ADDR'];
-        $requestCount = get_transient("rate_limit_{$ip}");
+        $key = $this->getRateLimitKey($request);
+        $requestCount = get_transient($key);
 
         if ($requestCount === false) {
-            set_transient("rate_limit_{$ip}", 1, $this->windowTime);
+            set_transient($key, 1, $this->windowTime);
         } else {
             if ($requestCount >= $this->maxRequests) {
+                do_action('plubo/rate_limit_exceeded', $request, $this->type);
                 return new \WP_REST_Response(['error' => 'Too many requests'], 429);
             }
-            set_transient("rate_limit_{$ip}", $requestCount + 1, $this->windowTime);
+            set_transient($key, $requestCount + 1, $this->windowTime);
         }
 
-        return $next($request);
+        $response = $next($request);
+
+        // Add rate limit headers
+        $response->header('X-RateLimit-Limit', $this->maxRequests);
+        $response->header('X-RateLimit-Remaining', max(0, $this->maxRequests - ($requestCount + 1)));
+        $response->header('X-RateLimit-Reset', time() + $this->windowTime);
+
+        return $response;
+    }
+
+    protected function getRateLimitKey($request)
+    {
+        switch ($this->type) {
+            case 'user':
+                $identifier = get_current_user_id() ?: $this->getClientIp();
+                break;
+            case 'endpoint':
+                $identifier = $request->get_route();
+                break;
+            case 'ip':
+            default:
+                $identifier = $this->getClientIp();
+        }
+        return $this->prefix . md5($identifier . $this->type);
+    }
+
+    protected function getClientIp()
+    {
+        foreach (array('HTTP_CLIENT_IP', 'HTTP_X_FORWARDED_FOR', 'HTTP_X_FORWARDED', 'HTTP_X_CLUSTER_CLIENT_IP', 'HTTP_FORWARDED_FOR', 'HTTP_FORWARDED', 'REMOTE_ADDR') as $key) {
+            if (array_key_exists($key, $_SERVER) === true) {
+                foreach (explode(',', $_SERVER[$key]) as $ip) {
+                    $ip = trim($ip);
+
+                    if (filter_var($ip, FILTER_VALIDATE_IP, FILTER_FLAG_NO_PRIV_RANGE | FILTER_FLAG_NO_RES_RANGE) !== false) {
+                        return $ip;
+                    }
+                }
+            }
+        }
     }
 }