Browse Source

Added std::ostream os in DataSink.

yhirose 5 years ago
parent
commit
25aa3ca982
2 changed files with 34 additions and 18 deletions
  1. 24 7
      httplib.h
  2. 10 11
      test/test.cc

+ 24 - 7
httplib.h

@@ -215,7 +215,8 @@ using MultipartFormDataMap = std::multimap<std::string, MultipartFormData>;
 
 class DataSink {
 public:
-  DataSink() = default;
+  DataSink() : os(&sb_), sb_(*this) {}
+
   DataSink(const DataSink &) = delete;
   DataSink &operator=(const DataSink &) = delete;
   DataSink(DataSink &&) = delete;
@@ -224,6 +225,24 @@ public:
   std::function<void(const char *data, size_t data_len)> write;
   std::function<void()> done;
   std::function<bool()> is_writable;
+  std::ostream os;
+
+private:
+  class data_sink_streambuf : public std::streambuf {
+  public:
+    data_sink_streambuf(DataSink &sink) : sink_(sink) {}
+
+  protected:
+    std::streamsize xsputn(const char *s, std::streamsize n) {
+      sink_.write(s, static_cast<size_t>(n));
+      return n;
+    }
+
+  private:
+    DataSink &sink_;
+  };
+
+  data_sink_streambuf sb_;
 };
 
 using ContentProvider =
@@ -2084,22 +2103,20 @@ inline ssize_t write_content(Stream &strm, ContentProvider content_provider,
   size_t begin_offset = offset;
   size_t end_offset = offset + length;
 
-  ssize_t written_length = 0;
+  auto ok = true;
 
   DataSink data_sink;
   data_sink.write = [&](const char *d, size_t l) {
     offset += l;
-    written_length = strm.write(d, l);
-  };
-  data_sink.is_writable = [&](void) {
-    return strm.is_writable() && written_length >= 0;
+    if (strm.write(d, l) < 0) { ok = false; }
   };
+  data_sink.is_writable = [&](void) { return strm.is_writable() && ok; };
 
   while (offset < end_offset) {
     if (!content_provider(offset, end_offset - offset, data_sink)) {
       return -1;
     }
-    if (written_length < 0) { return written_length; }
+    if (!ok) { return -1; }
   }
 
   return static_cast<ssize_t>(offset - begin_offset);

+ 10 - 11
test/test.cc

@@ -875,9 +875,9 @@ protected:
                res.set_chunked_content_provider(
                    [](size_t /*offset*/, DataSink &sink) {
                      EXPECT_TRUE(sink.is_writable());
-                     sink.write("123", 3);
-                     sink.write("456", 3);
-                     sink.write("789", 3);
+                     sink.os << "123";
+                     sink.os << "456";
+                     sink.os << "789";
                      sink.done();
                      return true;
                    });
@@ -889,9 +889,9 @@ protected:
                    [i](size_t /*offset*/, DataSink &sink) {
                      EXPECT_TRUE(sink.is_writable());
                      switch (*i) {
-                     case 0: sink.write("123", 3); break;
-                     case 1: sink.write("456", 3); break;
-                     case 2: sink.write("789", 3); break;
+                     case 0: sink.os << "123"; break;
+                     case 1: sink.os << "456"; break;
+                     case 2: sink.os << "789"; break;
                      case 3: sink.done(); break;
                      }
                      (*i)++;
@@ -903,7 +903,7 @@ protected:
              [&](const Request & /*req*/, Response &res) {
                res.set_content_provider(
                    6, [](size_t offset, size_t /*length*/, DataSink &sink) {
-                     sink.write(offset < 3 ? "a" : "b", 1);
+                     sink.os << (offset < 3 ? "a" : "b");
                      return true;
                    });
              })
@@ -929,8 +929,7 @@ protected:
                    size_t(-1),
                    [](size_t /*offset*/, size_t /*length*/, DataSink &sink) {
                      EXPECT_TRUE(sink.is_writable());
-                     std::string data = "data_chunk";
-                     sink.write(data.data(), data.size());
+                     sink.os << "data_chunk";
                      return true;
                    });
              })
@@ -1898,7 +1897,7 @@ TEST_F(ServerTest, PutWithContentProvider) {
       "/put", 3,
       [](size_t /*offset*/, size_t /*length*/, DataSink &sink) {
         EXPECT_TRUE(sink.is_writable());
-        sink.write("PUT", 3);
+        sink.os << "PUT";
         return true;
       },
       "text/plain");
@@ -1926,7 +1925,7 @@ TEST_F(ServerTest, PutWithContentProviderWithGzip) {
       "/put", 3,
       [](size_t /*offset*/, size_t /*length*/, DataSink &sink) {
         EXPECT_TRUE(sink.is_writable());
-        sink.write("PUT", 3);
+        sink.os << "PUT";
         return true;
       },
       "text/plain");