Regularly check bitrate

GuacLive · Dec 15, 2019 · 7908799 · 7908799
1 parent ed7271d
commit 7908799
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 0 deletions.
diff --git a/misc/index.js b/misc/index.js
@@ -3,6 +3,7 @@ const axios = require('axios');
 const cron = require('node-cron')
 // eslint-disable-next-line import/no-unresolved
 const helpers = require('./utils/helpers');
+const Logger = require('../node_core_logger');
 const conf = require('./config');
 
 const IS_DEBUG = process.env.NODE_ENV === 'development';
@@ -31,6 +32,8 @@ const config = {
     api_key: conf.api_key,
     ignore_auth: !!IS_DEBUG,
     maxDataRate: conf.maxDataRate || 8000,
+    dataRateCheckInterval: conf.dataRateCheckInterval || 3,
+    dataRateCheckCount: conf.dataRateCheckCount || 5, 
     transcode: conf.transcode
   }
 };
@@ -105,6 +108,7 @@ nms.on('onMetaData', (id, metadata) => {
   console.log('onMetaData', id, metadata);
   let session = nms.getSession(id);
   if(metadata.videodatarate > config.misc.maxDataRate){
+    Logger.error('Bitrate too high', `${Math.round(Math.floor(metadata.videodatarate))}/${config.misc.maxDataRate} kbps (max).`);
     session.sendStatusMessage(
       session.publishStreamId,
       'error',

diff --git a/node_rtmp_session.js b/node_rtmp_session.js
@@ -171,6 +171,10 @@ class NodeRtmpSession {
     this.numPlayCache = 0;
 
     context.sessions.set(this.id, this);
+
+    this.totalBytes = 0;
+    this.exceedDataRateCount = 0;
+    this.dataRateCheck = null;
   }
 
   run() {
@@ -191,6 +195,14 @@ class NodeRtmpSession {
         this.pingInterval = null;
       }
 
+      if (this.dataRateCheck != null) {
+        clearInterval(this.dataRateCheck);
+        this.dataRateCheck = null;
+      }
+
+      this.totalBytes = 0;
+      this.exceedDataRateCount = 0;
+
       if (this.playStreamId > 0) {
         this.onDeleteStream({ streamId: this.playStreamId });
       }
@@ -1068,6 +1080,29 @@ class NodeRtmpSession {
           }
         }
         context.nodeEvent.emit("postPublish", this.id, this.publishStreamPath, this.publishArgs);
+
+        this.dataRateCheck = setInterval(() => {
+          const bytes = this.socket.bytesRead - this.totalBytes;
+          this.totalBytes += bytes;
+          const dataRate = bytes / this.config.misc.dataRateCheckInterval / 125; // Kbps
+
+          if (dataRate > this.config.misc.maxDataRate + 10000) {
+            this.exceedDataRateCount++;
+          } else {
+            this.exceedDataRateCount = 0;
+          }
+
+          if (this.exceedDataRateCount >= this.config.misc.dataRateCheckCount) {
+            Logger.error('Bitrate too high', `${Math.round(Math.floor(metadata.videodatarate))}/${config.misc.maxDataRate} kbps (max).`);
+            this.sendStatusMessage(
+              this.publishStreamId,
+              'error',
+              'NetStream.Publish.Rejected',
+              `Bitrate too high, ${Math.round(Math.floor(dataRate))}/${config.misc.maxDataRate} kbps (max).`
+            );
+            this.reject();
+          }
+        }, this.config.misc.dataRateCheckInterval * 1000);
       }
       if (context.publishers.has(this.publishStreamPath)) {
         let publisherId = context.publishers.get(this.publishStreamPath);