refine bandwidth server-side, use bandwidth sample and kbps limit service

trunk/src/app/srs_app_bandwidth.cpp

@@ -36,6 +36,17 @@ using namespace std;
 #include <srs_core_autofree.hpp>
 #include <srs_kernel_utility.hpp>
 #include <srs_app_utility.hpp>
+#include <srs_app_kbps.hpp>
+
+SrsBandwidthSample::SrsBandwidthSample()
+{
+    duration_ms = 3000;
+    interval_ms = actual_duration_ms = bytes = 0;
+}
+
+SrsBandwidthSample::~SrsBandwidthSample()
+{
+}
 
 SrsBandwidth::SrsBandwidth()
 {
@@ -47,7 +58,7 @@ SrsBandwidth::~SrsBandwidth()
 {
 }
 
-int SrsBandwidth::bandwidth_check(SrsRtmpServer* rtmp, SrsRequest* req, string local_ip)
+int SrsBandwidth::bandwidth_check(SrsRtmpServer* rtmp, ISrsProtocolStatistic* io_stat, SrsRequest* req, string local_ip)
 {
     int ret = ERROR_SUCCESS;
     
@@ -68,7 +79,7 @@ int SrsBandwidth::bandwidth_check(SrsRtmpServer* rtmp, SrsRequest* req, string l
     }
     
     // shared global last check time,
-    // to avoid attach by bandwidth check,
+    // to prevent bandwidth check attack,
     // if client request check in the window(specifeid by interval),
     // directly reject the request.
     static int64_t last_check_time = 0;
@@ -94,24 +105,22 @@ int SrsBandwidth::bandwidth_check(SrsRtmpServer* rtmp, SrsRequest* req, string l
         return ret;
     }
     
-    return do_bandwidth_check();
+    // create a limit object.
+    SrsKbps kbps;
+    kbps.set_io(io_stat, io_stat);
+
+    int limit_kbps = _srs_config->get_bw_check_limit_kbps(_req->vhost);
+    SrsKbpsLimit limit(&kbps, limit_kbps);
+    
+    return do_bandwidth_check(&limit);
 }
 
-int SrsBandwidth::do_bandwidth_check()
+int SrsBandwidth::do_bandwidth_check(SrsKbpsLimit* limit)
 {
     int ret = ERROR_SUCCESS;
 
-    int play_duration_ms        = 3000;
-    int play_interval_ms        = 0;
-    int play_actual_duration_ms = 0;
-    int play_bytes              = 0;
-
-    int publish_duration_ms        = 3000;
-    int publish_interval_ms        = 0;
-    int publish_actual_duration_ms = 0;
-    int publish_bytes              = 0;
-    
-    int limit_kbps = _srs_config->get_bw_check_limit_kbps(_req->vhost);
+    SrsBandwidthSample play_sample;
+    SrsBandwidthSample publish_sample;
 
     int64_t start_time = srs_get_system_time_ms();
     

trunk/src/app/srs_app_bandwidth.hpp

@@ -35,6 +35,36 @@ CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
 class SrsRequest;
 class SrsRtmpServer;
+class SrsKbpsLimit;
+class ISrsProtocolStatistic;
+
+/**
+* bandwidth check/test sample.
+*/
+class SrsBandwidthSample
+{
+public:
+    /**
+    * the plan, how long to do the test, in ms,
+    * if exceed the duration, abort the test.
+    */
+    int duration_ms;
+    /**
+    * the plan, interval for each check/test packet, in ms
+    */
+    int interval_ms;
+    /**
+    * the actual test duration, in ms.
+    */
+    int actual_duration_ms;
+    /**
+    * the actual test bytes
+    */
+    int bytes;
+public:
+    SrsBandwidthSample();
+    virtual ~SrsBandwidthSample();
+};
 
 /**
 * bandwidth test agent which provides the interfaces for bandwidth check.
@@ -84,15 +114,17 @@ public:
     /**
     * do the bandwidth check.
     * @param rtmp, server RTMP protocol object, send/recv RTMP packet to/from client.
+    * @param io_stat, the underlayer io statistic, provides send/recv bytes count.
     * @param req, client request object, specifies the request info from client.
     * @param local_ip, the ip of server which client connected at
     */
-    virtual int bandwidth_check(SrsRtmpServer* rtmp, SrsRequest* req, std::string local_ip);
+    virtual int bandwidth_check(SrsRtmpServer* rtmp, ISrsProtocolStatistic* io_stat, SrsRequest* req, std::string local_ip);
 private:
     /**
     * used to process band width check from client.
+    * @param limit, the bandwidth limit object, to slowdown if exceed the kbps.
     */
-    virtual int do_bandwidth_check();
+    virtual int do_bandwidth_check(SrsKbpsLimit* limit);
     virtual int check_play(int duration_ms, int interval_ms, int& actual_duration_ms, int& play_bytes, int max_play_kbps);
     virtual int check_publish(int duration_ms, int interval_ms, int& actual_duration_ms, int& publish_bytes, int max_pub_kbps);
 };

trunk/src/app/srs_app_kbps.cpp

@@ -28,6 +28,8 @@ CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 #include <srs_protocol_io.hpp>
 #include <srs_kernel_utility.hpp>
 
+#define _SRS_BANDWIDTH_LIMIT_INTERVAL_MS 100
+
 SrsKbpsSample::SrsKbpsSample()
 {
     bytes = time = 0;
@@ -244,3 +246,29 @@ void SrsKbps::sample()
     os.sample();
 }
 
+SrsKbpsLimit::SrsKbpsLimit(SrsKbps* kbps, int limit_kbps)
+{
+    _kbps = kbps;
+    _limit_kbps = limit_kbps;
+}
+
+SrsKbpsLimit::~SrsKbpsLimit()
+{
+}
+
+void SrsKbpsLimit::recv_limit()
+{
+    while (_kbps->get_recv_kbps() > _limit_kbps) {
+        _kbps->sample();
+        st_usleep(_SRS_BANDWIDTH_LIMIT_INTERVAL_MS * 1000);
+    }
+}
+
+void SrsKbpsLimit::send_limit()
+{
+    while (_kbps->get_send_kbps() > _limit_kbps) {
+        _kbps->sample();
+        st_usleep(_SRS_BANDWIDTH_LIMIT_INTERVAL_MS * 1000);
+    }
+}
+

trunk/src/app/srs_app_kbps.hpp

@@ -198,4 +198,26 @@ public:
     virtual void sample();
 };
 
+/**
+* the kbps limit, if exceed the kbps, slow down.
+*/
+class SrsKbpsLimit
+{
+private:
+    int _limit_kbps;
+    SrsKbps* _kbps;
+public:
+    SrsKbpsLimit(SrsKbps* kbps, int limit_kbps);
+    virtual ~SrsKbpsLimit();
+public:
+    /**
+    * limit the recv bandwidth.
+    */
+    virtual void recv_limit();
+    /**
+    * limit the send bandwidth.
+    */
+    virtual void send_limit();
+};
+
 #endif

trunk/src/app/srs_app_rtmp_conn.cpp

@@ -216,7 +216,7 @@ int SrsRtmpConn::service_cycle()
     
     // do bandwidth test if connect to the vhost which is for bandwidth check.
     if (_srs_config->get_bw_check_enabled(req->vhost)) {
-        return bandwidth->bandwidth_check(rtmp, req, local_ip);
+        return bandwidth->bandwidth_check(rtmp, io, req, local_ip);
     }
     
     if ((ret = rtmp->response_connect_app(req, local_ip.c_str())) != ERROR_SUCCESS) {