diff --git a/dlib/sockstreambuf/sockstreambuf.cpp b/dlib/sockstreambuf/sockstreambuf.cpp index e328e42593..fb651a2d88 100644 --- a/dlib/sockstreambuf/sockstreambuf.cpp +++ b/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,7 +118,7 @@ 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 @@ -126,9 +126,9 @@ namespace dlib } // 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(*gptr()); } diff --git a/dlib/sockstreambuf/sockstreambuf.h b/dlib/sockstreambuf/sockstreambuf.h index f5b450e784..e0f587344d 100644 --- a/dlib/sockstreambuf/sockstreambuf.h +++ b/dlib/sockstreambuf/sockstreambuf.h @@ -4,7 +4,10 @@ #define DLIB_SOCKStREAMBUF_Hh_ #include +#include +#include #include +#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::max(), + "out_buffer_size must be in the range (0, INT_MAX]"); + DLIB_ASSERT(0 < in_buffer_size && in_buffer_size <= std::numeric_limits::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(out_buffer_size)]); + in_buffer.reset(new char[static_cast(in_buffer_size)]); init(); } sockstreambuf ( - const std::unique_ptr& con_ - ) : - con(*con_), - out_buffer(0), - in_buffer(0), - autoflush(false) - { - init(); - } + const std::unique_ptr& 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(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 out_buffer; + std::unique_ptr in_buffer; bool autoflush; }; diff --git a/dlib/sockstreambuf/sockstreambuf_abstract.h b/dlib/sockstreambuf/sockstreambuf_abstract.h index 12be84193a..2f9c49d5fb 100644 --- a/dlib/sockstreambuf/sockstreambuf_abstract.h +++ b/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::max() + - in_buffer_size > 0 + - in_buffer_size <= std::numeric_limits::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& con + const std::unique_ptr& 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::max() + - in_buffer_size > 0 + - in_buffer_size <= std::numeric_limits::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 !*/ diff --git a/dlib/test/sockstreambuf.cpp b/dlib/test/sockstreambuf.cpp index 519feb2b51..b89f1bd0ac 100644 --- a/dlib/test/sockstreambuf.cpp +++ b/dlib/test/sockstreambuf.cpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include @@ -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 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 con(raw_con); + if (accept_result != 0) + return; + + write_result = con->write(input.data(), input.size()); + if (write_result != static_cast(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 con(raw_con); + + vector 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(input.size())); + DLIB_TEST(read_result > 0); + DLIB_TEST(output_stream_good); + DLIB_TEST(received_output == output); + + const vector 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_custom_buffer_test(); dlog << LINFO << "testing sockstreambuf_unbuffered"; sockstreambuf_test(); }