Add: Limit Upload and Download speed

config.yaml

@@ -13,3 +13,6 @@ PASV:  # PASV 模式，指定该模式下可以使用的端口范围
   #   或者在网关采用动态 IP 的情况下有用
   force_passive_ip: "" # 若使用程序自动解析的 ip 地址，则填为空字符串 ""
 
+SERVICE:
+  max_upload_speed: 409600  # 上传速度, 400k Bytes/s
+  max_download_speed: 409600  # 下载速度, 400k Bytes/s

src/command_handle.cpp

@@ -12,6 +12,8 @@
 #include "record_lock.h"
 #include <unistd.h>
 #include <crypt.h>
+#include "speed_barrier.h"
+#include "configure.h"
 
 CLCommandHandle::CLCommandHandle(int command_fd, int read_pipe_fd) :
         m_pipe_fd(read_pipe_fd), m_cmd_fd(command_fd), m_b_stop(false),
@@ -303,12 +305,15 @@ void CLCommandHandle::do_retr() {   /* 下载文件 */
             file_size -= m_resume_point; /* 需要传送的字节的数目 */
         }
 
+        CLSpeedBarrier download_speed_barrier(configure::MAX_DOWNLOAD_SPEED);   /* 限速器 */
         while (file_size > 0) { /* 传输文件数据 */
             int sent_size = sendfile(data_connection->get_fd(), fd, nullptr, BYTES_PEER_TRANSFER); /* 每次发送一点 */
             if (sent_size == -1) {
+                utility::debug_info("Sent file failed");
                 break;
             }
             file_size -= sent_size;
+            download_speed_barrier.limit_speed(sent_size);  /* 限速 */
         }
     }
     if (file_size == 0) {
@@ -359,6 +364,7 @@ void CLCommandHandle::do_stor() {
 
         /* 读取数据,写入本地文件 */
         char buf[1024 * 1024] = {0};
+        CLSpeedBarrier upload_speed_barrier(configure::MAX_UPLOAD_SPEED);   /* 限速器 */
         while (true) {
             int res = read(data_connection->get_fd(), buf, sizeof(buf));
             if (res == -1) {
@@ -375,6 +381,7 @@ void CLCommandHandle::do_stor() {
                 state = ipc_utility::WriteError;
                 break;
             }
+            upload_speed_barrier.limit_speed(res);  /* 限速 */
         }
     }
     tcp::close_fd(fd); /* 关闭文件 */

src/command_handle.h

@@ -128,7 +128,7 @@ class CLCommandHandle {
     EMMode m_data_type; /* 传输的数据类型 */
 
     std::string m_file_name; /* 用来记录需要重命名的文件名 */
-    static const int BYTES_PEER_TRANSFER = 1024 * 1024; /* 文件下载时每次传输的数据量 */
+    static const int BYTES_PEER_TRANSFER = 1024 * 100; /* 文件下载时每次传输的数据量：100k Bytes/s */
     long long m_resume_point;   /* 断点续传点 */
 private:
     uid_t m_uid;    /* 用户 id  */

src/configure.cpp

@@ -21,6 +21,9 @@ namespace configure {
     int PASV_PORT_LOW;
     int PASV_PORT_HIGH;
 
+    int MAX_UPLOAD_SPEED; /* 最大上传速度 Bytes/s */
+    int MAX_DOWNLOAD_SPEED; /* 最大下载速度 Bytes/s */
+
     YAML::Node get_node(const YAML::Node &node, const std::string &node_name) { /* 获取节点 */
         if (!node[node_name.c_str()].IsDefined()) {
             std::cerr << "Configure information:" << node_name << " isn't configured correctly" << std::endl;
@@ -87,6 +90,15 @@ namespace configure {
             get_val(pasv_config, "force_passive_ip", FORCE_PASSIVE_SERVER_IP);
         }
         utility::debug_info(std::string("Force Server IP Address: ") + FORCE_PASSIVE_SERVER_IP);
+
+        auto service_config = get_node(node, "SERVICE");
+        get_val(service_config, "max_upload_speed", MAX_UPLOAD_SPEED);
+        get_val(service_config, "max_download_speed", MAX_DOWNLOAD_SPEED);
+
+        utility::debug_info(std::string("max upload speed: ") + std::to_string(MAX_UPLOAD_SPEED));
+        utility::debug_info(std::string("max download speed: ") + std::to_string(MAX_DOWNLOAD_SPEED));
+        assert(MAX_UPLOAD_SPEED > 0);
+        assert(MAX_DOWNLOAD_SPEED > 0);
     }
 
 }

src/configure.h

@@ -19,6 +19,9 @@ namespace configure {
     extern int PASV_PORT_LOW;
     extern int PASV_PORT_HIGH;
 
+    extern int MAX_UPLOAD_SPEED;
+    extern int MAX_DOWNLOAD_SPEED;
+
     void parse_config_file();   /* 从配置文件中解析配置信息 */
 
     YAML::Node get_node(const YAML::Node &node, const std::string &node_name);

src/speed_barrier.cpp

@@ -0,0 +1,33 @@
+//
+// Created by MLee on 2020/1/8.
+//
+
+#include "common.h"
+#include "speed_barrier.h"
+#include "utility.h"
+
+void CLSpeedBarrier::limit_speed(size_t transferred_size) {
+    auto time_now = std::chrono::system_clock::now();
+    std::chrono::duration<double> elapsed = time_now - m_start_time;
+    int64_t elapsed_micro_seconds = std::chrono::duration_cast<std::chrono::microseconds>(elapsed).count();
+    double elapsed_seconds = elapsed.count();
+
+    auto speed_now = static_cast<int64_t >(transferred_size / elapsed_seconds);
+    if (speed_now > m_max_speed) {
+        int64_t speed_diff = speed_now - m_max_speed;
+        int64_t need_sleep_micro_seconds = (speed_diff / static_cast<double >(m_max_speed)) * elapsed_micro_seconds;
+
+        /* 睡眠 */
+        struct timespec timeSpan{0};
+        timeSpan.tv_sec = need_sleep_micro_seconds / k_micro_sec_per_sec;
+        // 1微秒等于1000纳秒
+        timeSpan.tv_nsec = (need_sleep_micro_seconds - timeSpan.tv_sec * k_micro_sec_per_sec) * 1000; /* 纳秒级别的精度 */
+
+        int res;
+        do {
+            res = nanosleep(&timeSpan, &timeSpan);
+        } while (res == -1 && errno == EINTR); /* 可能被信号中断,所以要用do while格式 */
+    }
+
+    m_start_time = std::chrono::system_clock::now();    /* 重新计时 */
+}

src/speed_barrier.h

@@ -0,0 +1,26 @@
+//
+// Created by MLee on 2020/1/8.
+//
+
+#ifndef MINIFTPD_SPEED_BARRIER_H
+#define MINIFTPD_SPEED_BARRIER_H
+
+#include <chrono>
+
+class CLSpeedBarrier {
+public:
+    explicit CLSpeedBarrier(int64_t max_speed) : m_max_speed(max_speed) {
+        m_start_time = std::chrono::system_clock::now();
+    }
+
+    void limit_speed(size_t transferred_size);
+
+private:
+    static const int k_micro_sec_per_sec = 1000 * 1000; /* 1秒等于1000000微秒 */
+
+private:
+    int64_t m_max_speed;
+    std::chrono::time_point<std::chrono::system_clock> m_start_time;
+};
+
+#endif //MINIFTPD_SPEED_BARRIER_H