feat(ratelimit): Add Attributes support to rate limit middleware

Signed-off-by: Joas Schilling <coding@schilljs.com>

1 files changed, 170 insertions, 103 deletions

diff --git a/tests/lib/AppFramework/Middleware/Security/RateLimitingMiddlewareTest.php b/tests/lib/AppFramework/Middleware/Security/RateLimitingMiddlewareTest.php
index 38d01950f6a..47d479b18a1 100644
--- a/tests/lib/AppFramework/Middleware/Security/RateLimitingMiddlewareTest.php
+++ b/tests/lib/AppFramework/Middleware/Security/RateLimitingMiddlewareTest.php
@@ -1,7 +1,13 @@
 <?php
+
+declare(strict_types=1);
+
 /**
+ * @copyright Copyright (c) 2023 Joas Schilling <coding@schilljs.com>
  * @copyright Copyright (c) 2017 Lukas Reschke <lukas@statuscode.ch>
  *
+ * @author Joas Schilling <coding@schilljs.com>
+ *
  * @license GNU AGPL version 3 or any later version
  *
  * This program is free software: you can redistribute it and/or modify
@@ -26,34 +32,59 @@ use OC\AppFramework\Utility\ControllerMethodReflector;
 use OC\Security\RateLimiting\Exception\RateLimitExceededException;
 use OC\Security\RateLimiting\Limiter;
 use OCP\AppFramework\Controller;
+use OCP\AppFramework\Http\Attribute\AnonRateLimit;
+use OCP\AppFramework\Http\Attribute\UserRateLimit;
 use OCP\AppFramework\Http\DataResponse;
 use OCP\AppFramework\Http\TemplateResponse;
 use OCP\IRequest;
 use OCP\IUser;
 use OCP\IUserSession;
+use PHPUnit\Framework\MockObject\MockObject;
 use Test\TestCase;
 
+class TestRateLimitController extends Controller {
+	/**
+	 * @UserRateThrottle(limit=20, period=200)
+	 * @AnonRateThrottle(limit=10, period=100)
+	 */
+	public function testMethodWithAnnotation() {
+	}
+
+	/**
+	 * @AnonRateThrottle(limit=10, period=100)
+	 */
+	public function testMethodWithAnnotationFallback() {
+	}
+
+	public function testMethodWithoutAnnotation() {
+	}
+
+	#[UserRateLimit(limit: 20, period: 200)]
+	#[AnonRateLimit(limit: 10, period: 100)]
+	public function testMethodWithAttributes() {
+	}
+
+	#[AnonRateLimit(limit: 10, period: 100)]
+	public function testMethodWithAttributesFallback() {
+	}
+}
+
 /**
  * @group DB
  */
 class RateLimitingMiddlewareTest extends TestCase {
-	/** @var IRequest|\PHPUnit\Framework\MockObject\MockObject */
-	private $request;
-	/** @var IUserSession|\PHPUnit\Framework\MockObject\MockObject */
-	private $userSession;
-	/** @var ControllerMethodReflector|\PHPUnit\Framework\MockObject\MockObject */
-	private $reflector;
-	/** @var Limiter|\PHPUnit\Framework\MockObject\MockObject */
-	private $limiter;
-	/** @var RateLimitingMiddleware */
-	private $rateLimitingMiddleware;
+	private IRequest|MockObject $request;
+	private IUserSession|MockObject $userSession;
+	private ControllerMethodReflector $reflector;
+	private Limiter|MockObject $limiter;
+	private RateLimitingMiddleware $rateLimitingMiddleware;
 
 	protected function setUp(): void {
 		parent::setUp();
 
 		$this->request = $this->createMock(IRequest::class);
 		$this->userSession = $this->createMock(IUserSession::class);
-		$this->reflector = $this->createMock(ControllerMethodReflector::class);
+		$this->reflector = new ControllerMethodReflector();
 		$this->limiter = $this->createMock(Limiter::class);
 
 		$this->rateLimitingMiddleware = new RateLimitingMiddleware(
@@ -64,23 +95,25 @@ class RateLimitingMiddlewareTest extends TestCase {
 		);
 	}
 
-	public function testBeforeControllerWithoutAnnotation() {
-		$this->reflector
-			->expects($this->exactly(4))
-			->method('getAnnotationParameter')
-			->withConsecutive(
-				['AnonRateThrottle', 'limit'],
-				['AnonRateThrottle', 'period'],
-				['UserRateThrottle', 'limit'],
-				['UserRateThrottle', 'period']
-			)
-			->willReturnMap([
-				['AnonRateThrottle', 'limit', ''],
-				['AnonRateThrottle', 'period', ''],
-				['UserRateThrottle', 'limit', ''],
-				['UserRateThrottle', 'period', ''],
-			]);
+	public function testBeforeControllerWithoutAnnotationForAnon(): void {
+		$this->limiter
+			->expects($this->never())
+			->method('registerUserRequest');
+		$this->limiter
+			->expects($this->never())
+			->method('registerAnonRequest');
+
+		$this->userSession->expects($this->once())
+			->method('isLoggedIn')
+			->willReturn(false);
+
+		/** @var TestRateLimitController|MockObject $controller */
+		$controller = $this->createMock(TestRateLimitController::class);
+		$this->reflector->reflect($controller, 'testMethodWithoutAnnotation');
+		$this->rateLimitingMiddleware->beforeController($controller, 'testMethodWithoutAnnotation');
+	}
 
+	public function testBeforeControllerWithoutAnnotationForLoggedIn(): void {
 		$this->limiter
 			->expects($this->never())
 			->method('registerUserRequest');
@@ -88,34 +121,79 @@ class RateLimitingMiddlewareTest extends TestCase {
 			->expects($this->never())
 			->method('registerAnonRequest');
 
-		/** @var Controller|\PHPUnit\Framework\MockObject\MockObject $controller */
-		$controller = $this->createMock(Controller::class);
-		$this->rateLimitingMiddleware->beforeController($controller, 'testMethod');
+		$this->userSession->expects($this->once())
+			->method('isLoggedIn')
+			->willReturn(true);
+
+		/** @var TestRateLimitController|MockObject $controller */
+		$controller = $this->createMock(TestRateLimitController::class);
+		$this->reflector->reflect($controller, 'testMethodWithoutAnnotation');
+		$this->rateLimitingMiddleware->beforeController($controller, 'testMethodWithoutAnnotation');
 	}
 
-	public function testBeforeControllerForAnon() {
-		/** @var Controller|\PHPUnit\Framework\MockObject\MockObject $controller */
-		$controller = $this->createMock(Controller::class);
+	public function testBeforeControllerForAnon(): void {
+		$controller = new TestRateLimitController('test', $this->request);
+
+		$this->request
+			->method('getRemoteAddress')
+			->willReturn('127.0.0.1');
+
+		$this->userSession
+			->expects($this->once())
+			->method('isLoggedIn')
+			->willReturn(false);
+
+		$this->limiter
+			->expects($this->never())
+			->method('registerUserRequest');
+		$this->limiter
+			->expects($this->once())
+			->method('registerAnonRequest')
+			->with(get_class($controller) . '::testMethodWithAnnotation', '10', '100', '127.0.0.1');
+
+		$this->reflector->reflect($controller, 'testMethodWithAnnotation');
+		$this->rateLimitingMiddleware->beforeController($controller, 'testMethodWithAnnotation');
+	}
+
+	public function testBeforeControllerForLoggedIn(): void {
+		$controller = new TestRateLimitController('test', $this->request);
+		/** @var IUser|MockObject $user */
+		$user = $this->createMock(IUser::class);
+
+		$this->userSession
+			->expects($this->once())
+			->method('isLoggedIn')
+			->willReturn(true);
+		$this->userSession
+			->expects($this->once())
+			->method('getUser')
+			->willReturn($user);
+
+		$this->limiter
+			->expects($this->never())
+			->method('registerAnonRequest');
+		$this->limiter
+			->expects($this->once())
+			->method('registerUserRequest')
+			->with(get_class($controller) . '::testMethodWithAnnotation', '20', '200', $user);
+
+
+		$this->reflector->reflect($controller, 'testMethodWithAnnotation');
+		$this->rateLimitingMiddleware->beforeController($controller, 'testMethodWithAnnotation');
+	}
+
+	public function testBeforeControllerAnonWithFallback(): void {
+		$controller = new TestRateLimitController('test', $this->request);
 		$this->request
 			->expects($this->once())
 			->method('getRemoteAddress')
 			->willReturn('127.0.0.1');
 
-		$this->reflector
-			->expects($this->exactly(4))
-			->method('getAnnotationParameter')
-			->withConsecutive(
-				['AnonRateThrottle', 'limit'],
-				['AnonRateThrottle', 'period'],
-				['UserRateThrottle', 'limit'],
-				['UserRateThrottle', 'period']
-			)
-			->willReturnMap([
-				['AnonRateThrottle', 'limit', '100'],
-				['AnonRateThrottle', 'period', '10'],
-				['UserRateThrottle', 'limit', ''],
-				['UserRateThrottle', 'period', ''],
-			]);
+		$this->userSession
+			->expects($this->once())
+			->method('isLoggedIn')
+			->willReturn(true);
+
 
 		$this->limiter
 			->expects($this->never())
@@ -123,16 +201,39 @@ class RateLimitingMiddlewareTest extends TestCase {
 		$this->limiter
 			->expects($this->once())
 			->method('registerAnonRequest')
-			->with(get_class($controller) . '::testMethod', '100', '10', '127.0.0.1');
+			->with(get_class($controller) . '::testMethodWithAnnotationFallback', '10', '100', '127.0.0.1');
+
+		$this->reflector->reflect($controller, 'testMethodWithAnnotationFallback');
+		$this->rateLimitingMiddleware->beforeController($controller, 'testMethodWithAnnotationFallback');
+	}
+
+	public function testBeforeControllerAttributesForAnon(): void {
+		$controller = new TestRateLimitController('test', $this->request);
+
+		$this->request
+			->method('getRemoteAddress')
+			->willReturn('127.0.0.1');
 
+		$this->userSession
+			->expects($this->once())
+			->method('isLoggedIn')
+			->willReturn(false);
 
-		$this->rateLimitingMiddleware->beforeController($controller, 'testMethod');
+		$this->limiter
+			->expects($this->never())
+			->method('registerUserRequest');
+		$this->limiter
+			->expects($this->once())
+			->method('registerAnonRequest')
+			->with(get_class($controller) . '::testMethodWithAttributes', '10', '100', '127.0.0.1');
+
+		$this->reflector->reflect($controller, 'testMethodWithAttributes');
+		$this->rateLimitingMiddleware->beforeController($controller, 'testMethodWithAttributes');
 	}
 
-	public function testBeforeControllerForLoggedIn() {
-		/** @var Controller|\PHPUnit\Framework\MockObject\MockObject $controller */
-		$controller = $this->createMock(Controller::class);
-		/** @var IUser|\PHPUnit\Framework\MockObject\MockObject $user */
+	public function testBeforeControllerAttributesForLoggedIn(): void {
+		$controller = new TestRateLimitController('test', $this->request);
+		/** @var IUser|MockObject $user */
 		$user = $this->createMock(IUser::class);
 
 		$this->userSession
@@ -144,37 +245,21 @@ class RateLimitingMiddlewareTest extends TestCase {
 			->method('getUser')
 			->willReturn($user);
 
-		$this->reflector
-			->expects($this->exactly(4))
-			->method('getAnnotationParameter')
-			->withConsecutive(
-				['AnonRateThrottle', 'limit'],
-				['AnonRateThrottle', 'period'],
-				['UserRateThrottle', 'limit'],
-				['UserRateThrottle', 'period']
-			)
-			->willReturnMap([
-				['AnonRateThrottle', 'limit', ''],
-				['AnonRateThrottle', 'period', ''],
-				['UserRateThrottle', 'limit', '100'],
-				['UserRateThrottle', 'period', '10'],
-			]);
-
 		$this->limiter
 			->expects($this->never())
 			->method('registerAnonRequest');
 		$this->limiter
 			->expects($this->once())
 			->method('registerUserRequest')
-			->with(get_class($controller) . '::testMethod', '100', '10', $user);
+			->with(get_class($controller) . '::testMethodWithAttributes', '20', '200', $user);
 
 
-		$this->rateLimitingMiddleware->beforeController($controller, 'testMethod');
+		$this->reflector->reflect($controller, 'testMethodWithAttributes');
+		$this->rateLimitingMiddleware->beforeController($controller, 'testMethodWithAttributes');
 	}
 
-	public function testBeforeControllerAnonWithFallback() {
-		/** @var Controller|\PHPUnit\Framework\MockObject\MockObject $controller */
-		$controller = $this->createMock(Controller::class);
+	public function testBeforeControllerAttributesAnonWithFallback(): void {
+		$controller = new TestRateLimitController('test', $this->request);
 		$this->request
 			->expects($this->once())
 			->method('getRemoteAddress')
@@ -183,23 +268,8 @@ class RateLimitingMiddlewareTest extends TestCase {
 		$this->userSession
 			->expects($this->once())
 			->method('isLoggedIn')
-			->willReturn(false);
+			->willReturn(true);
 
-		$this->reflector
-			->expects($this->exactly(4))
-			->method('getAnnotationParameter')
-			->withConsecutive(
-				['AnonRateThrottle', 'limit'],
-				['AnonRateThrottle', 'period'],
-				['UserRateThrottle', 'limit'],
-				['UserRateThrottle', 'period']
-			)
-			->willReturnMap([
-				['AnonRateThrottle', 'limit', '200'],
-				['AnonRateThrottle', 'period', '20'],
-				['UserRateThrottle', 'limit', '100'],
-				['UserRateThrottle', 'period', '10'],
-			]);
 
 		$this->limiter
 			->expects($this->never())
@@ -207,25 +277,23 @@ class RateLimitingMiddlewareTest extends TestCase {
 		$this->limiter
 			->expects($this->once())
 			->method('registerAnonRequest')
-			->with(get_class($controller) . '::testMethod', '200', '20', '127.0.0.1');
+			->with(get_class($controller) . '::testMethodWithAttributesFallback', '10', '100', '127.0.0.1');
 
-		$this->rateLimitingMiddleware->beforeController($controller, 'testMethod');
+		$this->reflector->reflect($controller, 'testMethodWithAttributesFallback');
+		$this->rateLimitingMiddleware->beforeController($controller, 'testMethodWithAttributesFallback');
 	}
 
-
-	public function testAfterExceptionWithOtherException() {
+	public function testAfterExceptionWithOtherException(): void {
 		$this->expectException(\Exception::class);
 		$this->expectExceptionMessage('My test exception');
 
-		/** @var Controller|\PHPUnit\Framework\MockObject\MockObject $controller */
-		$controller = $this->createMock(Controller::class);
+		$controller = new TestRateLimitController('test', $this->request);
 
 		$this->rateLimitingMiddleware->afterException($controller, 'testMethod', new \Exception('My test exception'));
 	}
 
-	public function testAfterExceptionWithJsonBody() {
-		/** @var Controller|\PHPUnit\Framework\MockObject\MockObject $controller */
-		$controller = $this->createMock(Controller::class);
+	public function testAfterExceptionWithJsonBody(): void {
+		$controller = new TestRateLimitController('test', $this->request);
 		$this->request
 			->expects($this->once())
 			->method('getHeader')
@@ -238,9 +306,8 @@ class RateLimitingMiddlewareTest extends TestCase {
 		$this->assertEquals($expected, $result);
 	}
 
-	public function testAfterExceptionWithHtmlBody() {
-		/** @var Controller|\PHPUnit\Framework\MockObject\MockObject $controller */
-		$controller = $this->createMock(Controller::class);
+	public function testAfterExceptionWithHtmlBody(): void {
+		$controller = new TestRateLimitController('test', $this->request);
 		$this->request
 			->expects($this->once())
 			->method('getHeader')