Allow setting sockstreambuf buffer size and max putback

dlib/sockstreambuf/sockstreambuf.cpp

@@ -107,7 +107,7 @@ namespace dlib
         }
 
         // copy the putback characters into the putback end of the in_buffer
-        std::memmove(in_buffer+(max_putback-num_put_back), gptr()-num_put_back, num_put_back);
+        std::memmove(in_buffer.get()+(max_putback-num_put_back), gptr()-num_put_back, num_put_back);
 
         if (flushes_output_on_read())
         {
@@ -118,17 +118,17 @@ namespace dlib
             }
         }
 
-        int num = con.read(in_buffer+max_putback, in_buffer_size-max_putback);
+        int num = con.read(in_buffer.get()+max_putback, in_buffer_size-max_putback);
         if (num <= 0)
         {
             // an error occurred or the connection is over which is EOF
             return EOF;
         }
 
         // reset in_buffer pointers
-        setg (in_buffer+(max_putback-num_put_back),
-              in_buffer+max_putback,
-              in_buffer+max_putback+num);
+        setg (in_buffer.get()+(max_putback-num_put_back),
+              in_buffer.get()+max_putback,
+              in_buffer.get()+max_putback+num);
 
         return static_cast<unsigned char>(*gptr());
     }

dlib/sockstreambuf/sockstreambuf.h

@@ -4,7 +4,10 @@
 #define DLIB_SOCKStREAMBUF_Hh_
 
 #include <iosfwd>
+#include <limits>
+#include <memory>
 #include <streambuf>
+#include "../assert.h"
 #include "../sockets.h"
 #include "sockstreambuf_abstract.h"
 #include "sockstreambuf_unbuffered.h"
@@ -36,33 +39,46 @@ namespace dlib
         typedef sockstreambuf kernel_2a;
 
         sockstreambuf (
-            connection* con_
+            connection* con_,
+            std::streamsize out_buffer_size_ = 10000,
+            std::streamsize in_buffer_size_ = 10000,
+            std::streamsize max_putback_ = 4
         ) :
             con(*con_),
-            out_buffer(0),
-            in_buffer(0),
+            out_buffer_size(out_buffer_size_),
+            in_buffer_size(in_buffer_size_),
+            max_putback(max_putback_),
             autoflush(false)
         {
+            DLIB_ASSERT(0 < out_buffer_size && out_buffer_size <= std::numeric_limits<int>::max(),
+                "out_buffer_size must be in the range (0, INT_MAX]");
+            DLIB_ASSERT(0 < in_buffer_size && in_buffer_size <= std::numeric_limits<int>::max(),
+                "in_buffer_size must be in the range (0, INT_MAX]");
+            DLIB_ASSERT(0 <= max_putback && max_putback < in_buffer_size,
+                "max_putback must be in the range [0, in_buffer_size)");
+
+            out_buffer.reset(new char[static_cast<size_t>(out_buffer_size)]);
+            in_buffer.reset(new char[static_cast<size_t>(in_buffer_size)]);
             init();
         }
 
         sockstreambuf (
-            const std::unique_ptr<connection>& con_
-        ) :
-            con(*con_),
-            out_buffer(0),
-            in_buffer(0),
-            autoflush(false)
-        {
-            init();
-        }
+            const std::unique_ptr<connection>& con_,
+            std::streamsize out_buffer_size_ = 10000,
+            std::streamsize in_buffer_size_ = 10000,
+            std::streamsize max_putback_ = 4
+        ) : sockstreambuf(
+                con_.get(),
+                out_buffer_size_,
+                in_buffer_size_,
+                max_putback_
+            )
+        {}
 
         virtual ~sockstreambuf (
         )
         {
             sync();
-            delete [] out_buffer;
-            delete [] in_buffer;
         }
 
         connection* get_connection (
@@ -88,27 +104,17 @@ namespace dlib
         void init (
         )
         {
-            try
-            {
-                out_buffer = new char[out_buffer_size];
-                in_buffer = new char[in_buffer_size];
-            }
-            catch (...)
-            {
-                if (out_buffer) delete [] out_buffer;
-                throw;
-            }
-            setp(out_buffer, out_buffer + (out_buffer_size-1));
-            setg(in_buffer+max_putback, 
-                 in_buffer+max_putback, 
-                 in_buffer+max_putback);
+            setp(out_buffer.get(), out_buffer.get() + (out_buffer_size-1));
+            setg(in_buffer.get()+max_putback,
+                 in_buffer.get()+max_putback,
+                 in_buffer.get()+max_putback);
         }
 
         int flush_out_buffer (
         )
         {
             int num = static_cast<int>(pptr()-pbase());
-            if (con.write(out_buffer,num) != num)
+            if (con.write(out_buffer.get(),num) != num)
             {
                 // the write was not successful so return EOF 
                 return EOF;
@@ -151,11 +157,11 @@ namespace dlib
 
         // member data
         connection&  con;
-        static const std::streamsize max_putback = 4;
-        static const std::streamsize out_buffer_size = 10000;
-        static const std::streamsize in_buffer_size = 10000;
-        char* out_buffer;
-        char* in_buffer;
+        const std::streamsize out_buffer_size;
+        const std::streamsize in_buffer_size;
+        const std::streamsize max_putback;
+        std::unique_ptr<char[]> out_buffer;
+        std::unique_ptr<char[]> in_buffer;
         bool autoflush;
 
     };

dlib/sockstreambuf/sockstreambuf_abstract.h

@@ -49,27 +49,51 @@ namespace dlib
         !*/
     public:
         sockstreambuf (
-            connection* con
+            connection* con,
+            std::streamsize out_buffer_size = 10000,
+            std::streamsize in_buffer_size = 10000,
+            std::streamsize max_putback = 4
         );
         /*!
             requires
                 - con == a valid connection object
+                - out_buffer_size > 0
+                - out_buffer_size <= std::numeric_limits<int>::max()
+                - in_buffer_size > 0
+                - in_buffer_size <= std::numeric_limits<int>::max()
+                - max_putback >= 0
+                - max_putback < in_buffer_size
             ensures
                 - *this will read from and write to con
                 - #flushes_output_on_read() == false
+                - out_buffer_size bytes are allocated for the internal output buffer
+                - in_buffer_size bytes are used to buffer input, with max_putback of those
+                  bytes reserved for putback characters
             throws
                 - std::bad_alloc
         !*/
 
         sockstreambuf (
-            const std::unique_ptr<connection>& con
+            const std::unique_ptr<connection>& con,
+            std::streamsize out_buffer_size = 10000,
+            std::streamsize in_buffer_size = 10000,
+            std::streamsize max_putback = 4
         );
         /*!
             requires
                 - con == a valid connection object
+                - out_buffer_size > 0
+                - out_buffer_size <= std::numeric_limits<int>::max()
+                - in_buffer_size > 0
+                - in_buffer_size <= std::numeric_limits<int>::max()
+                - max_putback >= 0
+                - max_putback < in_buffer_size
             ensures
                 - *this will read from and write to con
                 - #flushes_output_on_read() == false
+                - out_buffer_size bytes are allocated for the internal output buffer
+                - in_buffer_size bytes are used to buffer input, with max_putback of those
+                  bytes reserved for putback characters
             throws
                 - std::bad_alloc
         !*/

dlib/test/sockstreambuf.cpp

@@ -5,6 +5,7 @@
 #include <cstdlib>
 #include <sstream>
 #include <string>
+#include <thread>
 #include <vector>
 
 #include <ctime>
@@ -226,6 +227,97 @@ namespace
 
     }
 
+// ----------------------------------------------------------------------------------------
+
+    void sockstreambuf_custom_buffer_test (
+    )
+    {
+        listener* raw_list = 0;
+        DLIB_TEST(create_listener(raw_list, 0) == 0);
+        std::unique_ptr<listener> list(raw_list);
+
+        const string input = "abcdef";
+        const string output = "uvwxyz";
+        string received_output;
+        int accept_result = OTHER_ERROR;
+        long write_result = OTHER_ERROR;
+        long read_result = 0;
+
+        std::thread server([&]()
+        {
+            connection* raw_con = 0;
+            accept_result = list->accept(raw_con, 5000);
+            std::unique_ptr<connection> con(raw_con);
+            if (accept_result != 0)
+                return;
+
+            write_result = con->write(input.data(), input.size());
+            if (write_result != static_cast<long>(input.size()))
+                return;
+
+            while (received_output.size() < output.size())
+            {
+                char buf[10];
+                read_result = con->read(buf, output.size()-received_output.size(), 5000);
+                if (read_result <= 0)
+                    return;
+                received_output.append(buf, read_result);
+            }
+        });
+
+        connection* raw_con = 0;
+        const int connect_result = create_connection(
+            raw_con,
+            list->get_listening_port(),
+            "127.0.0.1"
+        );
+        std::unique_ptr<connection> con(raw_con);
+
+        vector<sockstreambuf::int_type> read_values;
+        bool output_stream_good = false;
+        if (connect_result == 0)
+        {
+            sockstreambuf buf(con, 1, 5, 3);
+
+            read_values.push_back(buf.sbumpc()); // a
+            read_values.push_back(buf.sbumpc()); // b
+            read_values.push_back(buf.sbumpc()); // c
+            read_values.push_back(buf.sbumpc()); // d
+            read_values.push_back(buf.sgetc());  // e, and refill the input buffer
+            read_values.push_back(buf.sungetc());
+            read_values.push_back(buf.sungetc());
+            read_values.push_back(buf.sungetc());
+            read_values.push_back(buf.sungetc());
+            read_values.push_back(buf.sbumpc());
+            read_values.push_back(buf.sbumpc());
+            read_values.push_back(buf.sbumpc());
+            read_values.push_back(buf.sbumpc());
+            read_values.push_back(buf.sbumpc());
+
+            ostream out(&buf);
+            for (char ch : output)
+                out.put(ch);
+            out.flush();
+            output_stream_good = out.good();
+        }
+
+        server.join();
+
+        DLIB_TEST(connect_result == 0);
+        DLIB_TEST(accept_result == 0);
+        DLIB_TEST(write_result == static_cast<long>(input.size()));
+        DLIB_TEST(read_result > 0);
+        DLIB_TEST(output_stream_good);
+        DLIB_TEST(received_output == output);
+
+        const vector<sockstreambuf::int_type> expected_values = {
+            'a', 'b', 'c', 'd', 'e', 'd', 'c', 'b',
+            sockstreambuf::traits_type::eof(),
+            'b', 'c', 'd', 'e', 'f'
+        };
+        DLIB_TEST(read_values == expected_values);
+    }
+
 // ----------------------------------------------------------------------------------------
 
 
@@ -243,6 +335,7 @@ namespace
         {
             dlog << LINFO << "testing sockstreambuf";
             sockstreambuf_test<sockstreambuf>();
+            sockstreambuf_custom_buffer_test();
             dlog << LINFO << "testing sockstreambuf_unbuffered";
             sockstreambuf_test<sockstreambuf_unbuffered>();
         }