Fix quota inconstistencies with lives

4 files changed, 67 insertions, 26 deletions

diff --git a/server/lib/live/live-manager.ts b/server/lib/live/live-manager.ts
index df2804a0e..e04ae9fef 100644
--- a/server/lib/live/live-manager.ts
+++ b/server/lib/live/live-manager.ts
@@ -331,6 +331,8 @@ class LiveManager {
     muxingSession.on('after-cleanup', ({ videoId }) => {
       this.muxingSessions.delete(sessionId)
 
+      LiveQuotaStore.Instance.removeLive(user.id, videoLive.id)
+
       muxingSession.destroy()
 
       return this.onAfterMuxingCleanup({ videoId, liveSession })
diff --git a/server/lib/user.ts b/server/lib/user.ts
index 310a3c30c..586fd0f12 100644
--- a/server/lib/user.ts
+++ b/server/lib/user.ts
@@ -173,7 +173,8 @@ async function getOriginalVideoFileTotalFromUser (user: MUserId) {
   // Don't use sequelize because we need to use a sub query
   const query = UserModel.generateUserQuotaBaseSQL({
     withSelect: true,
-    whereUserId: '$userId'
+    whereUserId: '$userId',
+    daily: false
   })
 
   const base = await UserModel.getTotalRawQuery(query, user.id)
@@ -187,7 +188,7 @@ async function getOriginalVideoFileTotalDailyFromUser (user: MUserId) {
   const query = UserModel.generateUserQuotaBaseSQL({
     withSelect: true,
     whereUserId: '$userId',
-    where: '"video"."createdAt" > now() - interval \'24 hours\''
+    daily: true
   })
 
   const base = await UserModel.getTotalRawQuery(query, user.id)
diff --git a/server/models/user/user.ts b/server/models/user/user.ts
index bcf56dfa1..85720abfb 100644
--- a/server/models/user/user.ts
+++ b/server/models/user/user.ts
@@ -72,10 +72,13 @@ import { VideoImportModel } from '../video/video-import'
 import { VideoLiveModel } from '../video/video-live'
 import { VideoPlaylistModel } from '../video/video-playlist'
 import { UserNotificationSettingModel } from './user-notification-setting'
+import { LiveQuotaStore } from '@server/lib/live'
+import { logger } from '@server/helpers/logger'
 
 enum ScopeNames {
   FOR_ME_API = 'FOR_ME_API',
   WITH_VIDEOCHANNELS = 'WITH_VIDEOCHANNELS',
+  WITH_QUOTA = 'WITH_QUOTA',
   WITH_STATS = 'WITH_STATS'
 }
 
@@ -153,7 +156,7 @@ enum ScopeNames {
       }
     ]
   },
-  [ScopeNames.WITH_STATS]: {
+  [ScopeNames.WITH_QUOTA]: {
     attributes: {
       include: [
         [
@@ -161,7 +164,8 @@ enum ScopeNames {
             '(' +
               UserModel.generateUserQuotaBaseSQL({
                 withSelect: false,
-                whereUserId: '"UserModel"."id"'
+                whereUserId: '"UserModel"."id"',
+                daily: false
               }) +
             ')'
           ),
@@ -170,6 +174,24 @@ enum ScopeNames {
         [
           literal(
             '(' +
+              UserModel.generateUserQuotaBaseSQL({
+                withSelect: false,
+                whereUserId: '"UserModel"."id"',
+                daily: true
+              }) +
+            ')'
+          ),
+          'videoQuotaUsedDaily'
+        ]
+      ]
+    }
+  },
+  [ScopeNames.WITH_STATS]: {
+    attributes: {
+      include: [
+        [
+          literal(
+            '(' +
               'SELECT COUNT("video"."id") ' +
               'FROM "video" ' +
               'INNER JOIN "videoChannel" ON "videoChannel"."id" = "video"."channelId" ' +
@@ -474,21 +496,6 @@ export class UserModel extends Model<Partial<AttributesOnly<UserModel>>> {
     }
 
     const query: FindOptions = {
-      attributes: {
-        include: [
-          [
-            literal(
-              '(' +
-                UserModel.generateUserQuotaBaseSQL({
-                  withSelect: false,
-                  whereUserId: '"UserModel"."id"'
-                }) +
-              ')'
-            ),
-            'videoQuotaUsed'
-          ]
-        ]
-      },
       offset: start,
       limit: count,
       order: getSort(sort),
@@ -497,7 +504,7 @@ export class UserModel extends Model<Partial<AttributesOnly<UserModel>>> {
 
     return Promise.all([
       UserModel.unscoped().count(query),
-      UserModel.findAll(query)
+      UserModel.scope([ 'defaultScope', ScopeNames.WITH_QUOTA ]).findAll(query)
     ]).then(([ total, data ]) => ({ total, data }))
   }
 
@@ -579,7 +586,10 @@ export class UserModel extends Model<Partial<AttributesOnly<UserModel>>> {
       ScopeNames.WITH_VIDEOCHANNELS
     ]
 
-    if (withStats) scopes.push(ScopeNames.WITH_STATS)
+    if (withStats) {
+      scopes.push(ScopeNames.WITH_QUOTA)
+      scopes.push(ScopeNames.WITH_STATS)
+    }
 
     return UserModel.scope(scopes).findByPk(id)
   }
@@ -760,10 +770,10 @@ export class UserModel extends Model<Partial<AttributesOnly<UserModel>>> {
   static generateUserQuotaBaseSQL (options: {
     whereUserId: '$userId' | '"UserModel"."id"'
     withSelect: boolean
-    where?: string
+    daily: boolean
   }) {
-    const andWhere = options.where
-      ? 'AND ' + options.where
+    const andWhere = options.daily === true
+      ? 'AND "video"."createdAt" > now() - interval \'24 hours\''
       : ''
 
     const videoChannelJoin = 'INNER JOIN "videoChannel" ON "videoChannel"."id" = "video"."channelId" ' +
@@ -904,12 +914,15 @@ export class UserModel extends Model<Partial<AttributesOnly<UserModel>>> {
 
       videoQuota: this.videoQuota,
       videoQuotaDaily: this.videoQuotaDaily,
+
       videoQuotaUsed: videoQuotaUsed !== undefined
-        ? parseInt(videoQuotaUsed + '', 10)
+        ? parseInt(videoQuotaUsed + '', 10) + LiveQuotaStore.Instance.getLiveQuotaOf(this.id)
         : undefined,
+
       videoQuotaUsedDaily: videoQuotaUsedDaily !== undefined
-        ? parseInt(videoQuotaUsedDaily + '', 10)
+        ? parseInt(videoQuotaUsedDaily + '', 10) + LiveQuotaStore.Instance.getLiveQuotaOf(this.id)
         : undefined,
+
       videosCount: videosCount !== undefined
         ? parseInt(videosCount + '', 10)
         : undefined,
diff --git a/server/tests/api/live/live-constraints.ts b/server/tests/api/live/live-constraints.ts
index cab76f425..cf43c262a 100644
--- a/server/tests/api/live/live-constraints.ts
+++ b/server/tests/api/live/live-constraints.ts
@@ -12,6 +12,7 @@ import {
   PeerTubeServer,
   setAccessTokensToServers,
   setDefaultVideoChannel,
+  stopFfmpeg,
   waitJobs,
   waitUntilLiveReplacedByReplayOnAllServers,
   waitUntilLiveWaitingOnAllServers
@@ -169,6 +170,30 @@ describe('Test live constraints', function () {
     await servers[0].live.runAndTestStreamError({ token: userAccessToken, videoId: userVideoLiveoId, shouldHaveError: false })
   })
 
+  it('Should have the same quota in admin and as a user', async function () {
+    this.timeout(120000)
+
+    const userVideoLiveoId = await createLiveWrapper({ replay: true, permanent: false })
+    const ffmpegCommand = await servers[0].live.sendRTMPStreamInVideo({ token: userAccessToken, videoId: userVideoLiveoId })
+
+    await servers[0].live.waitUntilPublished({ videoId: userVideoLiveoId })
+
+    await wait(3000)
+
+    const quotaUser = await servers[0].users.getMyQuotaUsed({ token: userAccessToken })
+
+    const { data } = await servers[0].users.list()
+    const quotaAdmin = data.find(u => u.username === 'user1')
+
+    expect(quotaUser.videoQuotaUsed).to.equal(quotaAdmin.videoQuotaUsed)
+    expect(quotaUser.videoQuotaUsedDaily).to.equal(quotaAdmin.videoQuotaUsedDaily)
+
+    expect(quotaUser.videoQuotaUsed).to.be.above(10)
+    expect(quotaUser.videoQuotaUsedDaily).to.be.above(10)
+
+    await stopFfmpeg(ffmpegCommand)
+  })
+
   it('Should have max duration limit', async function () {
     this.timeout(60000)