From 8be7a7d2190889e6539ec3d1821bfc28171ece4f Mon Sep 17 00:00:00 2001
From: KernelDeimos <eric.alex.dube@gmail.com>
Date: Thu, 25 Jul 2024 20:23:05 -0400
Subject: [PATCH] dev: add logical rate-limit to new call method

---
 src/backend/src/services/SUService.js         | 11 +++++++---
 .../src/services/drivers/DriverService.js     | 20 +++++++++++++++++++
 2 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/src/backend/src/services/SUService.js b/src/backend/src/services/SUService.js
index 94124532..30491b55 100644
--- a/src/backend/src/services/SUService.js
+++ b/src/backend/src/services/SUService.js
@@ -22,10 +22,15 @@ class SUService extends BaseService {
     async get_system_actor () {
         return this.sys_actor_;
     }
-    async sudo (callback) {
+    async sudo (actor, callback) {
+        if ( ! callback ) {
+            callback = actor;
+            actor = await this.sys_actor_;
+        }
+        actor = Actor.adapt(actor);
         return await Context.get().sub({
-            user: await this.sys_user_,
-            actor: await this.sys_actor_,
+            actor,
+            user: actor.type.user,
         }).arun(callback);
     }
 }
diff --git a/src/backend/src/services/drivers/DriverService.js b/src/backend/src/services/drivers/DriverService.js
index 354c05b7..5e9ef348 100644
--- a/src/backend/src/services/drivers/DriverService.js
+++ b/src/backend/src/services/drivers/DriverService.js
@@ -24,6 +24,7 @@ const BaseService = require("../BaseService");
 const { Driver } = require("../../definitions/Driver");
 const { PermissionUtil } = require("../auth/PermissionService");
 const { Invoker } = require("@heyputer/puter-js-common/src/libs/invoker");
+const { get_user } = require("../../helpers");
 
 /**
  * DriverService provides the functionality of Puter drivers.
@@ -261,6 +262,9 @@ class DriverService extends BaseService {
             );
         }
 
+        const policy_holder = await get_user(
+            { username: effective_policy.holder });
+
         // NOT FINAL: this will be handled by 'get_policies_for_option_'
         // when cascading monthly usage is implemented.
         const svc_systemData = this.services.get('system-data');
@@ -276,6 +280,22 @@ class DriverService extends BaseService {
             
         const invoker = Invoker.create({
             decorators: [
+                {
+                    name: 'enforce logical rate-limit',
+                    on_call: async args => {
+                        if ( ! effective_policy['rate-limit'] ) return args;
+                        const svc_su = this.services.get('su');
+                        const svc_rateLimit = this.services.get('rate-limit');
+                        await svc_su.sudo(policy_holder, async () => {
+                            await svc_rateLimit.check_and_increment(
+                                `V1:${service_name}:${iface}:${method}`,
+                                effective_policy['rate-limit'].max,
+                                effective_policy['rate-limit'].period,
+                            );
+                        });
+                        return args;
+                    },
+                },
                 {
                     name: 'add metadata',
                     on_return: async result => {