diff --git a/lib/base/netstring.cpp b/lib/base/netstring.cpp index 027bb09d3..46c4041a2 100644 --- a/lib/base/netstring.cpp +++ b/lib/base/netstring.cpp @@ -34,11 +34,13 @@ using namespace icinga; */ StreamReadStatus NetString::ReadStringFromStream(const Stream::Ptr& stream, String *str, StreamReadContext& context) { - if (stream->IsEof()) + if (context.Eof) return StatusEof; - if (context.MustRead && context.FillFromStream(stream) == 0) + if (context.MustRead && !context.FillFromStream(stream)) { + context.Eof = true; return StatusEof; + } size_t header_length = 0; diff --git a/lib/base/stream.cpp b/lib/base/stream.cpp index 7ec89a3ef..84ba19887 100644 --- a/lib/base/stream.cpp +++ b/lib/base/stream.cpp @@ -63,11 +63,13 @@ void Stream::WaitForData(void) StreamReadStatus Stream::ReadLine(String *line, StreamReadContext& context) { - if (IsEof()) + if (context.Eof) return StatusEof; if (context.MustRead) { if (!context.FillFromStream(this)) { + context.Eof = true; + *line = String(context.Buffer, &(context.Buffer[context.Size])); boost::algorithm::trim_right(*line); @@ -106,6 +108,8 @@ bool StreamReadContext::FillFromStream(const Stream::Ptr& stream) if (Wait && stream->SupportsWaiting()) stream->WaitForData(); + size_t count = 0; + do { Buffer = (char *)realloc(Buffer, Size + 4096); @@ -114,13 +118,14 @@ bool StreamReadContext::FillFromStream(const Stream::Ptr& stream) size_t rc = stream->Read(Buffer + Size, 4096, true); - if (rc == 0 && stream->IsEof()) - return false; - Size += rc; + count += rc; } while (stream->IsDataAvailable()); - return true; + if (count == 0 && stream->IsEof()) + return false; + else + return true; } void StreamReadContext::DropData(size_t count) diff --git a/lib/base/stream.hpp b/lib/base/stream.hpp index 0dfeeb0bc..fcd14cc95 100644 --- a/lib/base/stream.hpp +++ b/lib/base/stream.hpp @@ -39,7 +39,7 @@ enum ConnectionRole struct StreamReadContext { StreamReadContext(bool wait = true) - : Buffer(NULL), Size(0), MustRead(true), Wait(wait) + : Buffer(NULL), Size(0), MustRead(true), Eof(false), Wait(wait) { } ~StreamReadContext(void) @@ -53,6 +53,7 @@ struct StreamReadContext char *Buffer; size_t Size; bool MustRead; + bool Eof; bool Wait; }; diff --git a/test/base-fifo.cpp b/test/base-fifo.cpp index 875c45670..e3781a5fd 100644 --- a/test/base-fifo.cpp +++ b/test/base-fifo.cpp @@ -43,12 +43,12 @@ BOOST_AUTO_TEST_CASE(io) BOOST_CHECK(fifo->GetAvailableBytes() == 5); char buffer1[2]; - fifo->Read(buffer1, 2); + fifo->Read(buffer1, 2, true); BOOST_CHECK(memcmp(buffer1, "he", 2) == 0); BOOST_CHECK(fifo->GetAvailableBytes() == 3); char buffer2[5]; - size_t rc = fifo->Read(buffer2, 5); + size_t rc = fifo->Read(buffer2, 5, true); BOOST_CHECK(rc == 3); BOOST_CHECK(memcmp(buffer2, "llo", 3) == 0); BOOST_CHECK(fifo->GetAvailableBytes() == 0); diff --git a/test/base-netstring.cpp b/test/base-netstring.cpp index 1f0f4c089..b1c499c97 100644 --- a/test/base-netstring.cpp +++ b/test/base-netstring.cpp @@ -33,7 +33,7 @@ BOOST_AUTO_TEST_CASE(netstring) String s; StreamReadContext src; - BOOST_CHECK(NetString::ReadStringFromStream(fifo, &s, src)); + BOOST_CHECK(NetString::ReadStringFromStream(fifo, &s, src) == StatusNewItem); BOOST_CHECK(s == "hello"); fifo->Close(); diff --git a/test/base-stream.cpp b/test/base-stream.cpp index 902c946f2..cc697bf31 100644 --- a/test/base-stream.cpp +++ b/test/base-stream.cpp @@ -37,18 +37,20 @@ BOOST_AUTO_TEST_CASE(readline_stdio) StreamReadContext rlc; String line; - BOOST_CHECK(stdstream->ReadLine(&line, rlc)); + BOOST_CHECK(stdstream->ReadLine(&line, rlc) == StatusNewItem); BOOST_CHECK(line == "Hello"); - BOOST_CHECK(stdstream->ReadLine(&line, rlc)); + BOOST_CHECK(stdstream->ReadLine(&line, rlc) == StatusNewItem); BOOST_CHECK(line == "World"); - BOOST_CHECK(stdstream->ReadLine(&line, rlc)); + BOOST_CHECK(stdstream->ReadLine(&line, rlc) == StatusNewItem); BOOST_CHECK(line == ""); - BOOST_CHECK(stdstream->ReadLine(&line, rlc)); + BOOST_CHECK(stdstream->ReadLine(&line, rlc) == StatusNewItem); BOOST_CHECK(line == ""); + BOOST_CHECK(stdstream->ReadLine(&line, rlc) == StatusEof); + stdstream->Close(); }