From cd6f3e6bbe2496dbf0f2e0ae1f5e2625402c16b3 Mon Sep 17 00:00:00 2001
From: Hongbo Wu <hongbo@omnivore.app>
Date: Tue, 9 Apr 2024 18:13:38 +0800
Subject: [PATCH] enforce usage limits for the API

---
 packages/api/src/apollo.ts                    | 18 ++++++++++++-
 packages/api/src/entity/service_usage.ts      | 25 +++++++++++++++++++
 packages/api/src/services/service_usage.ts    | 22 ++++++++++++++++
 .../db/migrations/0172.do.service_usage.sql   |  8 +++---
 4 files changed, 69 insertions(+), 4 deletions(-)
 create mode 100644 packages/api/src/entity/service_usage.ts
 create mode 100644 packages/api/src/services/service_usage.ts

diff --git a/packages/api/src/apollo.ts b/packages/api/src/apollo.ts
index 38420daa1..0326ec95a 100644
--- a/packages/api/src/apollo.ts
+++ b/packages/api/src/apollo.ts
@@ -29,6 +29,7 @@ import { logger } from './utils/logger'
 import { ReadingProgressDataSource } from './datasources/reading_progress_data_source'
 import { createPrometheusExporterPlugin } from '@bmatei/apollo-prometheus-exporter'
 import { ApolloServerPlugin } from 'apollo-server-plugin-base'
+import { countDailyServiceUsage } from './services/service_usage'
 
 const signToken = promisify(jwt.sign)
 const pubsub = createPubSubClient()
@@ -115,10 +116,25 @@ export function makeApolloServer(app: Express): ApolloServer {
 
   // enforce usage limits for the API
   const usageLimitPlugin = (): ApolloServerPlugin<RequestContext> => {
+    // TODO: load the limit from the DB into memory when the server starts
+    // hardcode the limit for now
+    const MAX_SENT_EMAIL_PER_DAY = 3
+
     return {
       async requestDidStart(contextValue) {
         // get graphql query from the request
-        console.log(contextValue)
+        const query = contextValue.request.query
+        // get the user id from the claims
+        const userId = contextValue.context.claims?.uid
+        const action = 'replyToEmail'
+        if (userId && query?.includes(action)) {
+          // get the user's email sent count from the DB
+          const emailSentCount = await countDailyServiceUsage(userId, action)
+          if (emailSentCount >= MAX_SENT_EMAIL_PER_DAY) {
+            // if the user has reached the limit, throw an error
+            throw new Error('You have reached the daily email limit')
+          }
+        }
       },
     }
   }
diff --git a/packages/api/src/entity/service_usage.ts b/packages/api/src/entity/service_usage.ts
new file mode 100644
index 000000000..056a51fc0
--- /dev/null
+++ b/packages/api/src/entity/service_usage.ts
@@ -0,0 +1,25 @@
+import {
+  Column,
+  CreateDateColumn,
+  Entity,
+  JoinColumn,
+  ManyToOne,
+  PrimaryGeneratedColumn,
+} from 'typeorm'
+import { User } from './user'
+
+@Entity('service_usage')
+export class ServiceUsage {
+  @PrimaryGeneratedColumn('uuid')
+  id!: string
+
+  @ManyToOne(() => User)
+  @JoinColumn({ name: 'user_id' })
+  user!: User
+
+  @Column('varchar')
+  action!: string
+
+  @CreateDateColumn({ default: () => 'CURRENT_TIMESTAMP' })
+  createdAt!: Date
+}
diff --git a/packages/api/src/services/service_usage.ts b/packages/api/src/services/service_usage.ts
new file mode 100644
index 000000000..9d4ada1d8
--- /dev/null
+++ b/packages/api/src/services/service_usage.ts
@@ -0,0 +1,22 @@
+import { Between } from 'typeorm'
+import { ServiceUsage } from '../entity/service_usage'
+import { authTrx, getRepository } from '../repository'
+import { DateTime } from 'luxon'
+
+const repo = getRepository(ServiceUsage)
+
+export const countDailyServiceUsage = async (
+  userId: string,
+  action: string
+) => {
+  return authTrx((tx) =>
+    tx.withRepository(repo).countBy({
+      user: { id: userId },
+      action,
+      createdAt: Between(
+        DateTime.now().startOf('day').toJSDate(),
+        DateTime.now().endOf('day').toJSDate()
+      ),
+    })
+  )
+}
diff --git a/packages/db/migrations/0172.do.service_usage.sql b/packages/db/migrations/0172.do.service_usage.sql
index f0beeb316..fa02e6214 100755
--- a/packages/db/migrations/0172.do.service_usage.sql
+++ b/packages/db/migrations/0172.do.service_usage.sql
@@ -28,9 +28,11 @@ CREATE TABLE omnivore.service_usage (
 
 CREATE INDEX ON omnivore.service_usage (user_id);
 
-CREATE POLICY create_service_usage on omnivore.service_usage
-    FOR INSERT TO omnivore_user
-    WITH CHECK (true);
+ALTER TABLE omnivore.service_usage ENABLE ROW LEVEL SECURITY;
+
+CREATE POLICY service_usage_policy on omnivore.service_usage
+    USING (user_id = omnivore.get_current_user_id())
+    WITH CHECK (user_id = omnivore.get_current_user_id());
 
 GRANT SELECT, INSERT ON omnivore.service_usage TO omnivore_user;