From 560854a9611ebe2361bebd7cc60b11a7e76ef2f9 Mon Sep 17 00:00:00 2001
From: Kent <kenthuang@synology.com>
Date: Fri, 12 Apr 2024 11:28:21 +0800
Subject: [PATCH] Apply range header base on response status code (#1806)

* Enable ignoring range header to generate customized response

* Apply range header base on response status code
---
 httplib.h    |  6 +++---
 test/test.cc | 23 +++++++++++++++++++++++
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/httplib.h b/httplib.h
index 4c446d9..4e79f11 100644
--- a/httplib.h
+++ b/httplib.h
@@ -6521,7 +6521,7 @@ inline bool Server::dispatch_request(Request &req, Response &res,
 inline void Server::apply_ranges(const Request &req, Response &res,
                                  std::string &content_type,
                                  std::string &boundary) const {
-  if (req.ranges.size() > 1) {
+  if (req.ranges.size() > 1 && res.status == StatusCode::PartialContent_206) {
     auto it = res.headers.find("Content-Type");
     if (it != res.headers.end()) {
       content_type = it->second;
@@ -6539,7 +6539,7 @@ inline void Server::apply_ranges(const Request &req, Response &res,
   if (res.body.empty()) {
     if (res.content_length_ > 0) {
       size_t length = 0;
-      if (req.ranges.empty()) {
+      if (req.ranges.empty() || res.status != StatusCode::PartialContent_206) {
         length = res.content_length_;
       } else if (req.ranges.size() == 1) {
         auto offset_and_length = detail::get_range_offset_and_length(
@@ -6568,7 +6568,7 @@ inline void Server::apply_ranges(const Request &req, Response &res,
       }
     }
   } else {
-    if (req.ranges.empty()) {
+    if (req.ranges.empty() || res.status != StatusCode::PartialContent_206) {
       ;
     } else if (req.ranges.size() == 1) {
       auto offset_and_length =
diff --git a/test/test.cc b/test/test.cc
index 71023f3..eae7dae 100644
--- a/test/test.cc
+++ b/test/test.cc
@@ -1890,6 +1890,11 @@ protected:
              [&](const Request & /*req*/, Response &res) {
                res.set_content("abcdefg", "text/plain");
              })
+        .Get("/with-range-customized-response",
+             [&](const Request & /*req*/, Response &res) {
+              res.status = StatusCode::BadRequest_400;
+              res.set_content(JSON_DATA,  "application/json");
+             })
         .Post("/chunked",
               [&](const Request &req, Response & /*res*/) {
                 EXPECT_EQ(req.body, "dechunked post body");
@@ -3166,6 +3171,24 @@ TEST_F(ServerTest, GetWithRangeMultipartOffsetGreaterThanContent) {
   EXPECT_EQ(StatusCode::RangeNotSatisfiable_416, res->status);
 }
 
+TEST_F(ServerTest, GetWithRangeCustomizedResponse) {
+  auto res = cli_.Get("/with-range-customized-response", {{make_range_header({{1, 2}})}});
+  ASSERT_TRUE(res);
+  EXPECT_EQ(StatusCode::BadRequest_400, res->status);
+  EXPECT_EQ(true, res->has_header("Content-Length"));
+  EXPECT_EQ(false, res->has_header("Content-Range"));
+  EXPECT_EQ(JSON_DATA, res->body);
+}
+
+TEST_F(ServerTest, GetWithRangeMultipartCustomizedResponseMultipleRange) {
+  auto res = cli_.Get("/with-range-customized-response", {{make_range_header({{1, 2}, {4, 5}})}});
+  ASSERT_TRUE(res);
+  EXPECT_EQ(StatusCode::BadRequest_400, res->status);
+  EXPECT_EQ(true, res->has_header("Content-Length"));
+  EXPECT_EQ(false, res->has_header("Content-Range"));
+  EXPECT_EQ(JSON_DATA, res->body);
+}
+
 TEST_F(ServerTest, Issue1772) {
   auto res = cli_.Get("/issue1772", {{make_range_header({{1000, -1}})}});
   ASSERT_TRUE(res);