ZLMediaKit · dennis-zheng · Mar 20, 2024
diff --git a/conf/config.ini b/conf/config.ini
@@ -225,6 +225,12 @@ charSet=utf-8
 keepAliveSecond=30
 #http请求体最大字节数，如果post的body太大，则不适合缓存body在内存
 maxReqSize=40960
+#http header请求最大个数，大于0才校验
+maxReqHeaderNumber=0
+#http header请求最大字节数，大于0才校验
+maxReqHeaderSize=0
+#http body请求最大字节数，大于0才校验
+maxReqBodySize=0
 #404网页内容，用户可以自定义404网页
 #notFound=<html><head><title>404 Not Found</title></head><body bgcolor="white"><center><h1>您访问的资源不存在！</h1></center><hr><center>ZLMediaKit-4.0</center></body></html>
 #http服务器监听端口

diff --git a/src/Common/config.cpp b/src/Common/config.cpp
@@ -164,6 +164,9 @@ namespace Http {
 #define HTTP_FIELD "http."
 const string kSendBufSize = HTTP_FIELD "sendBufSize";
 const string kMaxReqSize = HTTP_FIELD "maxReqSize";
+const string kMaxReqHeaderNumber = HTTP_FIELD "maxReqHeaderNumber";
+const string kMaxReqHeaderSize = HTTP_FIELD "maxReqHeaderSize";
+const string kMaxReqBodySize = HTTP_FIELD "maxReqBodySize";
 const string kKeepAliveSecond = HTTP_FIELD "keepAliveSecond";
 const string kCharSet = HTTP_FIELD "charSet";
 const string kRootPath = HTTP_FIELD "rootPath";

diff --git a/src/Common/config.h b/src/Common/config.h
@@ -254,6 +254,12 @@ namespace Http {
 extern const std::string kSendBufSize;
 // http 最大请求字节数
 extern const std::string kMaxReqSize;
+// header最大请求个数
+extern const std::string kMaxReqHeaderNumber;
+// header最大请求字节数
+extern const std::string kMaxReqHeaderSize;
+// body最大请求字节数
+extern const std::string kMaxReqBodySize;
 // http keep-alive秒数
 extern const std::string kKeepAliveSecond;
 // http 字符编码

diff --git a/src/Http/HttpRequestSplitter.cpp b/src/Http/HttpRequestSplitter.cpp
@@ -65,6 +65,10 @@ void HttpRequestSplitter::input(const char *data,size_t len) {
         _content_len = onRecvHeader(header_ptr, header_size);
     }
 
+    if (_content_len == 0 && _remain_data_size > 0) {
+        onCheckHeader(ptr,_remain_data_size);
+    }
+
     if(_remain_data_size <= 0){
         //没有剩余数据，清空缓存
         _remain_data.clear();

diff --git a/src/Http/HttpRequestSplitter.h b/src/Http/HttpRequestSplitter.h
@@ -68,7 +68,7 @@ class HttpRequestSplitter {
      * @param data content分片或全部数据
      * @param len 数据长度
      */
-    virtual void onRecvContent(const char *data,size_t len) {};
+    virtual void onRecvContent(const char *data, size_t len) {};
 
     /**
      * 判断数据中是否有包尾
@@ -78,6 +78,13 @@ class HttpRequestSplitter {
      */
     virtual const char *onSearchPacketTail(const char *data, size_t len);
 
+    /**
+     * 判断请求头是否有效
+     * @param data 请求头数据
+     * @param len 请求头长度
+     */
+    virtual void onCheckHeader(const char *data, size_t len) {};
+
     /**
      * 设置content len
      */

diff --git a/src/Http/HttpSession.cpp b/src/Http/HttpSession.cpp
@@ -26,6 +26,9 @@ namespace mediakit {
 HttpSession::HttpSession(const Socket::Ptr &pSock) : Session(pSock) {
     //设置默认参数
     setMaxReqSize(0);
+    setMaxReqHeaderNumber(0);
+    setMaxReqHeaderSize(0);
+    setMaxReqBodySize(0);
     setTimeoutSec(0);
 }
 
@@ -67,6 +70,8 @@ ssize_t HttpSession::onRecvHeader(const char *header, size_t len) {
     CHECK(_parser.url()[0] == '/');
     _origin = _parser["Origin"];
 
+    onCheckHeader(len, _parser.getHeader().size());
+
     urlDecode(_parser);
     auto &cmd = _parser.method();
     auto it = s_func_map.find(cmd);
@@ -140,11 +145,26 @@ ssize_t HttpSession::onRecvHeader(const char *header, size_t len) {
 }
 
 void HttpSession::onRecvContent(const char *data, size_t len) {
+    if (_max_req_body_size > 0 && len > _max_req_body_size) {
+        WarnL << "Http body size is too huge: " << len << " > " << _max_req_body_size
+                << ", please set " << Http::kMaxReqBodySize << " in config.ini file.";
+        reset();
+        throw std::out_of_range("Http body size is too huge: " + to_string(len));  
+    }
+
     if (_on_recv_body && !_on_recv_body(data, len)) {
         _on_recv_body = nullptr;
     }
 }
 
+void HttpSession::onCheckHeader(const char *header, size_t len) {
+    size_t number = 0;
+    if (_max_req_header_number > 0) {
+        number = onSearchHeaderNumber(header, len);
+    }
+    onCheckHeader(len, number);
+}
+
 void HttpSession::onRecv(const Buffer::Ptr &pBuf) {
     _ticker.resetTime();
     input(pBuf->data(), pBuf->size());
@@ -182,6 +202,30 @@ void HttpSession::setMaxReqSize(size_t max_req_size) {
     setMaxCacheSize(max_req_size);
 }
 
+void HttpSession::setMaxReqHeaderNumber(size_t max_req_header_number) {
+    if (!max_req_header_number) {
+        GET_CONFIG(size_t, s_max_req_header_number, Http::kMaxReqHeaderNumber);
+        max_req_header_number = s_max_req_header_number;
+    }
+    _max_req_header_number = max_req_header_number;
+}
+
+void HttpSession::setMaxReqHeaderSize(size_t max_req_header_size) {
+    if (!max_req_header_size) {
+        GET_CONFIG(size_t, s_max_req_header_size, Http::kMaxReqHeaderSize);
+        max_req_header_size = s_max_req_header_size;
+    }
+    _max_req_header_size = max_req_header_size;
+}
+
+void HttpSession::setMaxReqBodySize(size_t max_req_body_size) {
+    if (!max_req_body_size) {
+        GET_CONFIG(size_t, s_max_req_body_size, Http::kMaxReqBodySize);
+        max_req_body_size = s_max_req_body_size;
+    }
+    _max_req_body_size = max_req_body_size;
+}
+
 void HttpSession::onManager() {
     if (_ticker.elapsedTime() > _keep_alive_sec * 1000) {
         //http超时
@@ -834,4 +878,32 @@ std::shared_ptr<FlvMuxer> HttpSession::getSharedPtr() {
     return dynamic_pointer_cast<FlvMuxer>(shared_from_this());
 }
 
+ssize_t HttpSession::onSearchHeaderNumber(const char *data, size_t len) {
+    ssize_t number = 0;
+    const char *ptr = data;
+    const char *pos = nullptr;
+    while (pos = strstr(ptr, "\r\n")) {
+        number++;
+        ptr = pos + 2;
+    }
+    return number;
+}
+
+void HttpSession::onCheckHeader(size_t len, size_t number)
+{
+    if (_max_req_header_size > 0 && len > _max_req_header_size) {
+        WarnL << "Http header size is too huge: " << len << " > " << _max_req_header_size
+                << ", please set " << Http::kMaxReqHeaderSize << " in config.ini file.";
+        reset();
+        throw std::out_of_range("http header size is invalid: " + to_string(len));
+
+    }
+    if (_max_req_header_number > 0 && number > _max_req_header_number) {
+        WarnL << "Http header number is too huge: " << len << " > " << _max_req_header_number
+                << ", please set " << Http::kMaxReqHeaderNumber << " in config.ini file.";
+        reset();
+        throw std::out_of_range("http header size is invalid: " + to_string(len));
+    }
+}
+
 } /* namespace mediakit */
diff --git a/src/Http/HttpSession.h b/src/Http/HttpSession.h
@@ -49,6 +49,9 @@ class HttpSession: public toolkit::Session,
     static std::string urlDecodeComponent(const std::string &str);
     void setTimeoutSec(size_t second);
     void setMaxReqSize(size_t max_req_size);
+    void setMaxReqHeaderNumber(size_t max_req_header_number);
+    void setMaxReqHeaderSize(size_t max_req_header_size);
+    void setMaxReqBodySize(size_t max_req_body_size);
 
 protected:
     //FlvMuxer override
@@ -57,8 +60,9 @@ class HttpSession: public toolkit::Session,
     std::shared_ptr<FlvMuxer> getSharedPtr() override;
 
     //HttpRequestSplitter override
-    ssize_t onRecvHeader(const char *data,size_t len) override;
-    void onRecvContent(const char *data,size_t len) override;
+    ssize_t onRecvHeader(const char *data, size_t len) override;
+    void onRecvContent(const char *data, size_t len) override;
+    void onCheckHeader(const char *data, size_t len) override;
 
     /**
      * 重载之用于处理不定长度的content
@@ -126,6 +130,9 @@ class HttpSession: public toolkit::Session,
     //设置socket标志
     void setSocketFlags();
 
+    ssize_t onSearchHeaderNumber(const char *data, size_t len);
+    void onCheckHeader(size_t len, size_t number);
+
 protected:
     MediaInfo _media_info;
 
@@ -136,6 +143,9 @@ class HttpSession: public toolkit::Session,
     size_t _keep_alive_sec = 0;
     //最大http请求字节大小
     size_t _max_req_size = 0;
+    size_t _max_req_header_number = 0;
+    size_t _max_req_header_size = 0;
+    size_t _max_req_body_size = 0;
     //消耗的总流量
     uint64_t _total_bytes_usage = 0;
     // http请求中的 Origin字段