feat: limit storage size (#106)

packages/be-gateway/src/caches/OrganizationCache.ts

@@ -0,0 +1,24 @@
+import { mdOrgGetOne } from "@shared/models"
+import { CKEY, getCache, setCache } from "../lib/redis"
+
+export default class OrganizationCache {
+  orgId: string
+  key: string[]
+  constructor(orgId: string) {
+    this.orgId = orgId
+    this.key = [CKEY.ORG_MAX_STORAGE_SIZE, orgId]
+  }
+
+  async getMaxStorageSize() {
+    const cached = await getCache(this.key)
+    if (cached) {
+      return cached
+    }
+
+    const { maxStorageSize } = await mdOrgGetOne(this.orgId)
+    await setCache(this.key, maxStorageSize)
+
+    return maxStorageSize
+  }
+
+}

packages/be-gateway/src/caches/StorageCache.ts

@@ -0,0 +1,21 @@
+import { CKEY, decrByCache, getCache, incrByCache } from "../lib/redis";
+
+export default class StorageCache {
+  key: string[]
+  constructor(orgId: string) {
+    this.key = [CKEY.ORG_STORAGE_SIZE, orgId]
+  }
+
+  async getTotalSize() {
+    const size = await getCache(this.key)
+    return parseInt(size, 10)
+  }
+
+  async incrSize(size: number) {
+    return await incrByCache(this.key, size)
+  }
+
+  async decrSize(size: number) {
+    return await decrByCache(this.key, size)
+  }
+}

packages/be-gateway/src/lib/redis.ts

@@ -25,7 +25,10 @@ export enum CKEY {
   TASK_QUERY = 'TASK_QUERY',
 
   // save user's favorite links
-  FAV_QUERY = 'FAV_QUERY'
+  FAV_QUERY = 'FAV_QUERY',
+
+  ORG_STORAGE_SIZE = 'ORG_STORAGE_SIZE',
+  ORG_MAX_STORAGE_SIZE = 'ORG_MAX_STORAGE_SIZE',
 }
 
 type CACHE_KEY = CKEY | (CKEY | string)[]
@@ -209,6 +212,15 @@ export const incrCache = async (key: CACHE_KEY) => {
   }
 }
 
+export const incrByCache = async (key: CACHE_KEY, val: number) => {
+  try {
+    return await redis.incrby(genKey(key), val)
+  } catch (error) {
+    console.log(error)
+    return null
+  }
+}
+
 export const decrCache = async (key: CACHE_KEY) => {
   try {
     return await redis.decr(genKey(key))
@@ -217,3 +229,12 @@ export const decrCache = async (key: CACHE_KEY) => {
     return null
   }
 }
+
+export const decrByCache = async (key: CACHE_KEY, val: number) => {
+  try {
+    return await redis.decrby(genKey(key), val)
+  } catch (error) {
+    console.log(error)
+    return null
+  }
+}

packages/be-gateway/src/routes/organization/index.ts

@@ -14,6 +14,7 @@ import {
 } from '@shared/models'
 import orgMembers from './members'
 import { CKEY, delCache, getJSONCache, setJSONCache } from '../../lib/redis'
+import { MAX_STORAGE_SIZE } from '../storage'
 
 const router = Router()
 
@@ -60,13 +61,10 @@ router.post('/org', async (req: AuthRequest, res) => {
     const { id } = req.authen
     const key = [CKEY.USER_ORGS, id]
 
-    console.log('called organization')
-    console.log(body)
-    console.log(req.authen)
-
     const result = await mdOrgAdd({
       name: body.name,
       desc: body.desc,
+      maxStorageSize: MAX_STORAGE_SIZE,
       cover: null,
       avatar: null,
       createdAt: new Date(),
@@ -77,8 +75,6 @@ router.post('/org', async (req: AuthRequest, res) => {
 
     delCache(key)
 
-    console.log('created', result)
-
     await mdOrgMemAdd({
       uid: id,
       status: InvitationStatus.ACCEPTED,

packages/be-gateway/src/routes/storage/index.ts

@@ -7,6 +7,8 @@ import {
   randomObjectKeyName
 } from '@be/storage'
 import {
+  mdOrgGetOne,
+  mdProjectGetOrgId,
   mdStorageAdd,
   mdStorageGet,
   mdStorageGetByOwner,
@@ -16,9 +18,31 @@ import { FileOwnerType, FileStorage } from '@prisma/client'
 import { AuthRequest } from '../../types'
 import { pmClient } from 'packages/shared-models/src/lib/_prisma'
 import { CKEY, findNDelCaches } from '../../lib/redis'
+import StorageCache from '../../caches/StorageCache'
 
 const router = Router()
 
+export const MAX_STORAGE_SIZE = 100 * 1024 * 1024 // 100Mb
+router.get('/current-storage-size', async (req, res) => {
+  const { orgId } = req.query as { orgId: string }
+  if (!orgId) {
+    return res.status(500).send('ORG_MUST_PROVIDED')
+  }
+
+  const storageCache = new StorageCache(orgId)
+  const totalSize = await storageCache.getTotalSize()
+
+  const { maxStorageSize } = await mdOrgGetOne(orgId)
+
+  res.json({
+    status: 200,
+    data: {
+      maximum: maxStorageSize || MAX_STORAGE_SIZE,
+      total: totalSize
+    }
+  })
+})
+
 router.post('/create-presigned-url', async (req, res, next) => {
   const { name, type, orgId, projectId } = req.body as {
     name: string
@@ -29,6 +53,23 @@ router.post('/create-presigned-url', async (req, res, next) => {
 
   const randName = `${orgId}/${projectId}/` + randomObjectKeyName(name)
   const presignedUrl = await createPresignedUrlWithClient(randName, type)
+  const { organizationId } = await mdProjectGetOrgId(projectId)
+  const storageCache = new StorageCache(organizationId)
+  const totalSize = await storageCache.getTotalSize()
+  const { maxStorageSize } = await mdOrgGetOne(organizationId)
+
+  console.log('totalSize', totalSize)
+  console.log('maxStorageSize', maxStorageSize)
+
+  if (maxStorageSize && totalSize > maxStorageSize) {
+    res.status(500).send('MAX_SIZE_STORAGE')
+    return
+  }
+
+  if (totalSize > MAX_STORAGE_SIZE) {
+    res.status(500).send('MAX_SIZE_STORAGE')
+    return
+  }
 
   res.status(200).json({
     data: {
@@ -41,74 +82,87 @@ router.post('/create-presigned-url', async (req, res, next) => {
 
 router.delete('/del-file', async (req: AuthRequest, res) => {
   const { id, projectId } = req.query as { id: string; projectId: string }
-  const key = [CKEY.TASK_QUERY, projectId]
-
-  pmClient
-    .$transaction(async tx => {
-      const result = await tx.fileStorage.findFirst({
-        where: {
-          id
-        }
-      })
+  try {
+    const key = [CKEY.TASK_QUERY, projectId]
 
-      const { id: fileId, owner, ownerType, keyName } = result
+    const { organizationId } = await mdProjectGetOrgId(projectId)
+    const storageCache = new StorageCache(organizationId)
 
-      if (ownerType === FileOwnerType.TASK) {
-        const task = await tx.task.findFirst({
+    pmClient
+      .$transaction(async tx => {
+        const result = await tx.fileStorage.findFirst({
           where: {
-            id: owner
+            id
           }
         })
 
-        const { fileIds } = task
+        const { id: fileId, owner, ownerType, keyName } = result
 
-        if (!fileIds.includes(fileId)) {
-          // return 'FILE_NOT_EXIST_IN_TASK'
-          throw new Error('FILE_NOT_EXIST_IN_TASK')
-        }
+        if (ownerType === FileOwnerType.TASK) {
+          const task = await tx.task.findFirst({
+            where: {
+              id: owner
+            }
+          })
 
-        task.fileIds = fileIds.filter(f => f !== fileId)
+          const { fileIds } = task
 
-        delete task.id
+          if (!fileIds.includes(fileId)) {
+            // return 'FILE_NOT_EXIST_IN_TASK'
+            throw new Error('FILE_NOT_EXIST_IN_TASK')
+          }
 
-        const promises = []
-        promises.push(
-          tx.fileStorage.delete({
-            where: { id: fileId }
-          })
-        )
+          task.fileIds = fileIds.filter(f => f !== fileId)
 
-        promises.push(
-          tx.task.update({
-            where: {
-              id: owner
-            },
-            data: task
-          })
-        )
+          delete task.id
+
+          const promises = []
+          promises.push(
+            tx.fileStorage.delete({
+              where: { id: fileId }
+            })
+          )
 
-        await Promise.all(promises)
-        await deleteObject(keyName)
-        await findNDelCaches(key)
+          promises.push(
+            tx.task.update({
+              where: {
+                id: owner
+              },
+              data: task
+            })
+          )
 
-        return {
-          deletedFileId: fileId,
-          remainFileIds: task.fileIds
+          await Promise.all(promises)
+          await deleteObject(keyName)
+          await findNDelCaches(key)
+
+          // decrease storage size
+          const file = await mdStorageGetOne(fileId)
+          if (file && file.size) {
+            storageCache.decrSize(file.size)
+          }
+
+          return {
+            deletedFileId: fileId,
+            remainFileIds: task.fileIds
+          }
         }
-      }
 
-      // FIXME: this case only occurs as user create a new file in drive directly
-      // so, if it is not belong to no one, delete it
-      return 'CANNOT_DELETE_NO_OWNER'
-    })
-    .then(message => {
-      console.log('delete file result: ', message)
-      res.json({ status: 200, data: message })
-    })
-    .catch(err => {
-      console.log('error delete file', err)
-      res.status(500).send(err)
-    })
+        // FIXME: this case only occurs as user create a new file in drive directly
+        // so, if it is not belong to no one, delete it
+        return 'CANNOT_DELETE_NO_OWNER'
+      })
+      .then(message => {
+        console.log('delete file result: ', message)
+        res.json({ status: 200, data: message })
+      })
+      .catch(err => {
+        console.log('error delete file', err)
+        res.status(500).send(err)
+      })
+  } catch (error) {
+    res.status(500).send(error)
+  }
 })
 
 router.get('/get-files', async (req: AuthRequest, res) => {
@@ -154,28 +208,37 @@ router.post('/save-to-drive', async (req: AuthRequest, res, next) => {
     parentId
   } = req.body as FileStorage
 
-  const result = await mdStorageAdd({
-    organizationId,
-    projectId,
-    name,
-    keyName,
-    type,
-    url,
-    size,
-    mimeType,
-    parentId: parentId || null,
-    isDeleted: false,
-    owner,
-    ownerType,
-    createdAt: new Date(),
-    createdBy: uid,
-    deletedAt: null,
-    deletedBy: null
-  })
+  try {
+    if (size) {
+      const storageCache = new StorageCache(organizationId)
+      await storageCache.incrSize(size)
+    }
 
-  res.status(200).json({
-    data: result
-  })
+    const result = await mdStorageAdd({
+      organizationId,
+      projectId,
+      name,
+      keyName,
+      type,
+      url,
+      size,
+      mimeType,
+      parentId: parentId || null,
+      isDeleted: false,
+      owner,
+      ownerType,
+      createdAt: new Date(),
+      createdBy: uid,
+      deletedAt: null,
+      deletedBy: null
+    })
+
+    res.status(200).json({
+      data: result
+    })
+  } catch (error) {
+    res.status(500).send(error)
+  }
 })
 
 router.get('/get-object-url', async (req, res, next) => {