3 files changed, 140 insertions, 24 deletions

diff --git a/server/lib/job-queue/handlers/video-import.ts b/server/lib/job-queue/handlers/video-import.ts
index 9b5f2bb2b..9210aec54 100644
--- a/server/lib/job-queue/handlers/video-import.ts
+++ b/server/lib/job-queue/handlers/video-import.ts
@@ -4,6 +4,7 @@ import { extname } from 'path'
 import { addOptimizeOrMergeAudioJob } from '@server/helpers/video'
 import { isPostImportVideoAccepted } from '@server/lib/moderation'
 import { Hooks } from '@server/lib/plugins/hooks'
+import { isAbleToUploadVideo } from '@server/lib/user'
 import { getVideoFilePath } from '@server/lib/video-paths'
 import { MVideoImportDefault, MVideoImportDefaultFiles, MVideoImportVideo } from '@server/types/models/video/video-import'
 import {
@@ -108,7 +109,7 @@ async function processFile (downloader: () => Promise<string>, videoImport: MVid
 
     // Get information about this video
     const stats = await stat(tempVideoPath)
-    const isAble = await videoImport.User.isAbleToUploadVideo({ size: stats.size })
+    const isAble = await isAbleToUploadVideo(videoImport.User.id, stats.size)
     if (isAble === false) {
       throw new Error('The user video quota is exceeded with this video to import.')
     }
diff --git a/server/lib/live-manager.ts b/server/lib/live-manager.ts
index 41176d197..3ff2434ff 100644
--- a/server/lib/live-manager.ts
+++ b/server/lib/live-manager.ts
@@ -2,24 +2,27 @@
 import { AsyncQueue, queue } from 'async'
 import * as chokidar from 'chokidar'
 import { FfmpegCommand } from 'fluent-ffmpeg'
-import { ensureDir } from 'fs-extra'
+import { ensureDir, stat } from 'fs-extra'
 import { basename } from 'path'
 import { computeResolutionsToTranscode, runLiveMuxing, runLiveTranscoding } from '@server/helpers/ffmpeg-utils'
 import { logger } from '@server/helpers/logger'
 import { CONFIG, registerConfigChangedHandler } from '@server/initializers/config'
-import { P2P_MEDIA_LOADER_PEER_VERSION, VIDEO_LIVE, WEBSERVER } from '@server/initializers/constants'
+import { MEMOIZE_TTL, P2P_MEDIA_LOADER_PEER_VERSION, VIDEO_LIVE, WEBSERVER } from '@server/initializers/constants'
+import { UserModel } from '@server/models/account/user'
 import { VideoModel } from '@server/models/video/video'
 import { VideoFileModel } from '@server/models/video/video-file'
 import { VideoLiveModel } from '@server/models/video/video-live'
 import { VideoStreamingPlaylistModel } from '@server/models/video/video-streaming-playlist'
-import { MStreamingPlaylist, MVideoLiveVideo } from '@server/types/models'
+import { MStreamingPlaylist, MUser, MUserId, MVideoLive, MVideoLiveVideo } from '@server/types/models'
 import { VideoState, VideoStreamingPlaylistType } from '@shared/models'
 import { federateVideoIfNeeded } from './activitypub/videos'
 import { buildSha256Segment } from './hls'
 import { JobQueue } from './job-queue'
 import { PeerTubeSocket } from './peertube-socket'
+import { isAbleToUploadVideo } from './user'
 import { getHLSDirectory } from './video-paths'
 
+import memoizee = require('memoizee')
 const NodeRtmpServer = require('node-media-server/node_rtmp_server')
 const context = require('node-media-server/node_core_ctx')
 const nodeMediaServerLogger = require('node-media-server/node_core_logger')
@@ -53,6 +56,11 @@ class LiveManager {
   private readonly transSessions = new Map<string, FfmpegCommand>()
   private readonly videoSessions = new Map<number, string>()
   private readonly segmentsSha256 = new Map<string, Map<string, string>>()
+  private readonly livesPerUser = new Map<number, { liveId: number, videoId: number, size: number }[]>()
+
+  private readonly isAbleToUploadVideoWithCache = memoizee((userId: number) => {
+    return isAbleToUploadVideo(userId, 1000)
+  }, { maxAge: MEMOIZE_TTL.LIVE_ABLE_TO_UPLOAD })
 
   private segmentsSha256Queue: AsyncQueue<SegmentSha256QueueParam>
   private rtmpServer: any
@@ -127,7 +135,7 @@ class LiveManager {
 
     this.abortSession(sessionId)
 
-    this.onEndTransmuxing(videoId)
+    this.onEndTransmuxing(videoId, true)
       .catch(err => logger.error('Cannot end transmuxing of video %d.', videoId, { err }))
   }
 
@@ -196,8 +204,18 @@ class LiveManager {
     originalResolution: number
   }) {
     const { sessionId, videoLive, playlist, streamPath, resolutionsEnabled, originalResolution } = options
+    const startStreamDateTime = new Date().getTime()
     const allResolutions = resolutionsEnabled.concat([ originalResolution ])
 
+    const user = await UserModel.loadByLiveId(videoLive.id)
+    if (!this.livesPerUser.has(user.id)) {
+      this.livesPerUser.set(user.id, [])
+    }
+
+    const currentUserLive = { liveId: videoLive.id, videoId: videoLive.videoId, size: 0 }
+    const livesOfUser = this.livesPerUser.get(user.id)
+    livesOfUser.push(currentUserLive)
+
     for (let i = 0; i < allResolutions.length; i++) {
       const resolution = allResolutions[i]
 
@@ -216,26 +234,47 @@ class LiveManager {
     const outPath = getHLSDirectory(videoLive.Video)
     await ensureDir(outPath)
 
+    const deleteSegments = videoLive.saveReplay === false
+
     const rtmpUrl = 'rtmp://127.0.0.1:' + config.rtmp.port + streamPath
     const ffmpegExec = CONFIG.LIVE.TRANSCODING.ENABLED
-      ? runLiveTranscoding(rtmpUrl, outPath, allResolutions)
-      : runLiveMuxing(rtmpUrl, outPath)
+      ? runLiveTranscoding(rtmpUrl, outPath, allResolutions, deleteSegments)
+      : runLiveMuxing(rtmpUrl, outPath, deleteSegments)
 
     logger.info('Running live muxing/transcoding.')
-
     this.transSessions.set(sessionId, ffmpegExec)
 
     const videoUUID = videoLive.Video.uuid
     const tsWatcher = chokidar.watch(outPath + '/*.ts')
 
-    const updateHandler = segmentPath => {
-      this.segmentsSha256Queue.push({ operation: 'update', segmentPath, videoUUID })
+    const updateSegment = segmentPath => this.segmentsSha256Queue.push({ operation: 'update', segmentPath, videoUUID })
+
+    const addHandler = segmentPath => {
+      updateSegment(segmentPath)
+
+      if (this.isDurationConstraintValid(startStreamDateTime) !== true) {
+        this.stopSessionOf(videoLive.videoId)
+      }
+
+      if (videoLive.saveReplay === true) {
+        stat(segmentPath)
+          .then(segmentStat => {
+            currentUserLive.size += segmentStat.size
+          })
+          .then(() => this.isQuotaConstraintValid(user, videoLive))
+          .then(quotaValid => {
+            if (quotaValid !== true) {
+              this.stopSessionOf(videoLive.videoId)
+            }
+          })
+          .catch(err => logger.error('Cannot stat %s or check quota of %d.', segmentPath, user.id, { err }))
+      }
     }
 
     const deleteHandler = segmentPath => this.segmentsSha256Queue.push({ operation: 'delete', segmentPath, videoUUID })
 
-    tsWatcher.on('add', p => updateHandler(p))
-    tsWatcher.on('change', p => updateHandler(p))
+    tsWatcher.on('add', p => addHandler(p))
+    tsWatcher.on('change', p => updateSegment(p))
     tsWatcher.on('unlink', p => deleteHandler(p))
 
     const masterWatcher = chokidar.watch(outPath + '/master.m3u8')
@@ -280,7 +319,14 @@ class LiveManager {
     ffmpegExec.on('end', () => onFFmpegEnded())
   }
 
-  private async onEndTransmuxing (videoId: number) {
+  getLiveQuotaUsedByUser (userId: number) {
+    const currentLives = this.livesPerUser.get(userId)
+    if (!currentLives) return 0
+
+    return currentLives.reduce((sum, obj) => sum + obj.size, 0)
+  }
+
+  private async onEndTransmuxing (videoId: number, cleanupNow = false) {
     try {
       const fullVideo = await VideoModel.loadAndPopulateAccountAndServerAndTags(videoId)
       if (!fullVideo) return
@@ -290,7 +336,7 @@ class LiveManager {
         payload: {
           videoId: fullVideo.id
         }
-      }, { delay: VIDEO_LIVE.CLEANUP_DELAY })
+      }, { delay: cleanupNow ? 0 : VIDEO_LIVE.CLEANUP_DELAY })
 
       // FIXME: use end
       fullVideo.state = VideoState.WAITING_FOR_LIVE
@@ -337,6 +383,23 @@ class LiveManager {
     filesMap.delete(segmentName)
   }
 
+  private isDurationConstraintValid (streamingStartTime: number) {
+    const maxDuration = CONFIG.LIVE.MAX_DURATION
+    // No limit
+    if (maxDuration === null) return true
+
+    const now = new Date().getTime()
+    const max = streamingStartTime + maxDuration
+
+    return now <= max
+  }
+
+  private async isQuotaConstraintValid (user: MUserId, live: MVideoLive) {
+    if (live.saveReplay !== true) return true
+
+    return this.isAbleToUploadVideoWithCache(user.id)
+  }
+
   static get Instance () {
     return this.instance || (this.instance = new this())
   }
diff --git a/server/lib/user.ts b/server/lib/user.ts
index aa14f0b54..d3338f329 100644
--- a/server/lib/user.ts
+++ b/server/lib/user.ts
@@ -1,20 +1,24 @@
+import { Transaction } from 'sequelize/types'
 import { v4 as uuidv4 } from 'uuid'
+import { UserModel } from '@server/models/account/user'
 import { ActivityPubActorType } from '../../shared/models/activitypub'
+import { UserNotificationSetting, UserNotificationSettingValue } from '../../shared/models/users'
 import { SERVER_ACTOR_NAME, WEBSERVER } from '../initializers/constants'
+import { sequelizeTypescript } from '../initializers/database'
 import { AccountModel } from '../models/account/account'
-import { buildActorInstance, setAsyncActorKeys } from './activitypub/actor'
-import { createLocalVideoChannel } from './video-channel'
-import { ActorModel } from '../models/activitypub/actor'
 import { UserNotificationSettingModel } from '../models/account/user-notification-setting'
-import { UserNotificationSetting, UserNotificationSettingValue } from '../../shared/models/users'
-import { createWatchLaterPlaylist } from './video-playlist'
-import { sequelizeTypescript } from '../initializers/database'
-import { Transaction } from 'sequelize/types'
-import { Redis } from './redis'
-import { Emailer } from './emailer'
+import { ActorModel } from '../models/activitypub/actor'
 import { MAccountDefault, MActorDefault, MChannelActor } from '../types/models'
 import { MUser, MUserDefault, MUserId } from '../types/models/user'
+import { buildActorInstance, setAsyncActorKeys } from './activitypub/actor'
 import { getAccountActivityPubUrl } from './activitypub/url'
+import { Emailer } from './emailer'
+import { LiveManager } from './live-manager'
+import { Redis } from './redis'
+import { createLocalVideoChannel } from './video-channel'
+import { createWatchLaterPlaylist } from './video-playlist'
+
+import memoizee = require('memoizee')
 
 type ChannelNames = { name: string, displayName: string }
 
@@ -116,13 +120,61 @@ async function sendVerifyUserEmail (user: MUser, isPendingEmail = false) {
   await Emailer.Instance.addVerifyEmailJob(username, email, url)
 }
 
+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'
+  })
+
+  const base = await UserModel.getTotalRawQuery(query, user.id)
+
+  return base + LiveManager.Instance.getLiveQuotaUsedByUser(user.id)
+}
+
+// Returns cumulative size of all video files uploaded in the last 24 hours.
+async function getOriginalVideoFileTotalDailyFromUser (user: MUserId) {
+  // Don't use sequelize because we need to use a sub query
+  const query = UserModel.generateUserQuotaBaseSQL({
+    withSelect: true,
+    whereUserId: '$userId',
+    where: '"video"."createdAt" > now() - interval \'24 hours\''
+  })
+
+  const base = await UserModel.getTotalRawQuery(query, user.id)
+
+  return base + LiveManager.Instance.getLiveQuotaUsedByUser(user.id)
+}
+
+async function isAbleToUploadVideo (userId: number, size: number) {
+  const user = await UserModel.loadById(userId)
+
+  if (user.videoQuota === -1 && user.videoQuotaDaily === -1) return Promise.resolve(true)
+
+  const [ totalBytes, totalBytesDaily ] = await Promise.all([
+    getOriginalVideoFileTotalFromUser(user.id),
+    getOriginalVideoFileTotalDailyFromUser(user.id)
+  ])
+
+  const uploadedTotal = size + totalBytes
+  const uploadedDaily = size + totalBytesDaily
+
+  if (user.videoQuotaDaily === -1) return uploadedTotal < user.videoQuota
+  if (user.videoQuota === -1) return uploadedDaily < user.videoQuotaDaily
+
+  return uploadedTotal < user.videoQuota && uploadedDaily < user.videoQuotaDaily
+}
+
 // ---------------------------------------------------------------------------
 
 export {
+  getOriginalVideoFileTotalFromUser,
+  getOriginalVideoFileTotalDailyFromUser,
   createApplicationActor,
   createUserAccountAndChannelAndPlaylist,
   createLocalAccountWithoutKeys,
-  sendVerifyUserEmail
+  sendVerifyUserEmail,
+  isAbleToUploadVideo
 }
 
 // ---------------------------------------------------------------------------