From 58bf08111ab4636efe5808e4752271a59b858e37 Mon Sep 17 00:00:00 2001
From: GareArc <chen4851@purdue.edu>
Date: Wed, 5 Feb 2025 22:59:43 -0500
Subject: [PATCH] feat: add rate limiter to report download api

---
 api/controllers/console/error.py |  6 ++++++
 api/services/billing_service.py  | 12 +++++++++++-
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/api/controllers/console/error.py b/api/controllers/console/error.py
index ee87138a44..bd4ae9dc7f 100644
--- a/api/controllers/console/error.py
+++ b/api/controllers/console/error.py
@@ -101,3 +101,9 @@ class AccountInFreezeError(BaseHTTPException):
         "This email account has been deleted within the past 30 days"
         "and is temporarily unavailable for new account registration."
     )
+
+
+class CompilanceRateLimitError(BaseHTTPException):
+    error_code = "compilance_rate_limit"
+    description = "Rate limit exceeded for downloading compliance report."
+    code = 429
diff --git a/api/services/billing_service.py b/api/services/billing_service.py
index 8964c90b04..d4b3deb90a 100644
--- a/api/services/billing_service.py
+++ b/api/services/billing_service.py
@@ -6,6 +6,7 @@ from tenacity import (retry, retry_if_exception_type, stop_before_delay,
                       wait_fixed)
 
 from extensions.ext_database import db
+from libs.helper import RateLimiter
 from models.account import TenantAccountJoin, TenantAccountRole
 
 
@@ -13,6 +14,8 @@ class BillingService:
     base_url = os.environ.get("BILLING_API_URL", "BILLING_API_URL")
     secret_key = os.environ.get("BILLING_API_SECRET_KEY", "BILLING_API_SECRET_KEY")
 
+    compliance_download_rate_limiter = RateLimiter("compliance_download_rate_limiter", 3, 60)
+
     @classmethod
     def get_info(cls, tenant_id: str):
         params = {"tenant_id": tenant_id}
@@ -107,6 +110,11 @@ class BillingService:
         ip: str,
         device_info: str,
     ):
+        limiter_key = f"{account_id}:{tenant_id}"
+        if cls.compliance_download_rate_limiter.is_rate_limited(limiter_key):
+            from controllers.console.error import CompilanceRateLimitError
+            raise CompilanceRateLimitError()
+
         json = {
             "doc_name": doc_name,
             "account_id": account_id,
@@ -114,4 +122,6 @@ class BillingService:
             "ip_address": ip,
             "device_info": device_info,
         }
-        return cls._send_request("POST", "/compliance/download", json=json)
+        res = cls._send_request("POST", "/compliance/download", json=json)
+        cls.compliance_download_rate_limiter.increment(limiter_key)
+        return res