diff --git a/wpiutil/src/main/native/cpp/raw_istream.cpp b/wpiutil/src/main/native/cpp/raw_istream.cpp index f5aa97eae3..bdfb2bb0ad 100644 --- a/wpiutil/src/main/native/cpp/raw_istream.cpp +++ b/wpiutil/src/main/native/cpp/raw_istream.cpp @@ -47,9 +47,6 @@ StringRef raw_istream::getline(SmallVectorImpl& buf, int maxLen) { return StringRef{buf.data(), buf.size()}; } -raw_mem_istream::raw_mem_istream(StringRef mem) - : raw_mem_istream(mem.data(), mem.size()) {} - void raw_mem_istream::close() {} size_t raw_mem_istream::in_avail() const { return m_left; } @@ -57,11 +54,12 @@ size_t raw_mem_istream::in_avail() const { return m_left; } void raw_mem_istream::read_impl(void* data, size_t len) { if (len > m_left) { error_detected(); - return; + len = m_left; } std::memcpy(data, m_cur, len); m_cur += len; m_left -= len; + set_read_count(len); } static int getFD(const Twine& Filename, std::error_code& EC) { @@ -106,8 +104,9 @@ void raw_fd_istream::close() { size_t raw_fd_istream::in_avail() const { return m_end - m_cur; } void raw_fd_istream::read_impl(void* data, size_t len) { - size_t left = m_end - m_cur; - if (left < len) { + char* cdata = static_cast(data); + size_t pos = 0; + while (static_cast(m_end - m_cur) < len) { // not enough data if (m_cur == m_end) { #ifdef _WIN32 @@ -117,17 +116,23 @@ void raw_fd_istream::read_impl(void* data, size_t len) { #endif if (count <= 0) { error_detected(); + set_read_count(pos); return; } m_cur = m_buf; m_end = m_buf + count; - return read_impl(data, len); + continue; } - std::memcpy(data, m_cur, left); - return read_impl(static_cast(data) + left, len - left); + size_t left = m_end - m_cur; + std::memcpy(&cdata[pos], m_cur, left); + m_cur += left; + pos += left; + len -= left; } - std::memcpy(data, m_cur, len); + std::memcpy(&cdata[pos], m_cur, len); m_cur += len; + pos += len; + set_read_count(pos); } diff --git a/wpiutil/src/main/native/cpp/raw_socket_istream.cpp b/wpiutil/src/main/native/cpp/raw_socket_istream.cpp index 67a491eb19..68c4dd6b97 100644 --- a/wpiutil/src/main/native/cpp/raw_socket_istream.cpp +++ b/wpiutil/src/main/native/cpp/raw_socket_istream.cpp @@ -20,10 +20,11 @@ void raw_socket_istream::read_impl(void* data, size_t len) { size_t count = m_stream.receive(&cdata[pos], len - pos, &err, m_timeout); if (count == 0) { error_detected(); - return; + break; } pos += count; } + set_read_count(pos); } void raw_socket_istream::close() { m_stream.close(); } diff --git a/wpiutil/src/main/native/cpp/sha1.cpp b/wpiutil/src/main/native/cpp/sha1.cpp index 5c861e9063..e36eb45b53 100644 --- a/wpiutil/src/main/native/cpp/sha1.cpp +++ b/wpiutil/src/main/native/cpp/sha1.cpp @@ -215,7 +215,7 @@ static void buffer_to_block(const unsigned char* buffer, SHA1::SHA1() { reset(digest, buf_size, transforms); } void SHA1::Update(StringRef s) { - raw_mem_istream is(s); + raw_mem_istream is(makeArrayRef(s.data(), s.size())); Update(is); } diff --git a/wpiutil/src/main/native/include/wpi/raw_istream.h b/wpiutil/src/main/native/include/wpi/raw_istream.h index 2db724de6f..1e6faf9871 100644 --- a/wpiutil/src/main/native/include/wpi/raw_istream.h +++ b/wpiutil/src/main/native/include/wpi/raw_istream.h @@ -8,10 +8,15 @@ #ifndef WPIUTIL_WPI_RAW_ISTREAM_H_ #define WPIUTIL_WPI_RAW_ISTREAM_H_ +#include + #include #include +#include #include +#include +#include "wpi/ArrayRef.h" #include "wpi/SmallVector.h" #include "wpi/StringRef.h" #include "wpi/Twine.h" @@ -47,7 +52,47 @@ class raw_istream { size_t readlen = std::min(in_avail(), len); if (readlen == 0) return 0; read_impl(data, readlen); - return readlen; + return m_read_count; + } + + raw_istream& readinto(SmallVectorImpl& buf, size_t len) { + size_t old_size = buf.size(); + buf.append(len, 0); + read_impl(&buf[old_size], len); + buf.resize(old_size + m_read_count); + return *this; + } + + raw_istream& readinto(SmallVectorImpl& buf, size_t len) { + size_t old_size = buf.size(); + buf.append(len, 0); + read_impl(&buf[old_size], len); + buf.resize(old_size + m_read_count); + return *this; + } + + raw_istream& readinto(std::vector& buf, size_t len) { + size_t old_size = buf.size(); + buf.insert(buf.end(), len, 0); + read_impl(&buf[old_size], len); + buf.resize(old_size + m_read_count); + return *this; + } + + raw_istream& readinto(std::vector& buf, size_t len) { + size_t old_size = buf.size(); + buf.insert(buf.end(), len, 0); + read_impl(&buf[old_size], len); + buf.resize(old_size + m_read_count); + return *this; + } + + raw_istream& readinto(std::string& buf, size_t len) { + size_t old_size = buf.size(); + buf.insert(buf.end(), len, 0); + read_impl(&buf[old_size], len); + buf.resize(old_size + m_read_count); + return *this; } // Read a line from an input stream (up to a maximum length). @@ -59,8 +104,15 @@ class raw_istream { StringRef getline(SmallVectorImpl& buf, int maxLen); virtual void close() = 0; + + // Number of bytes available to read without potentially blocking. + // Note this can return zero even if there are bytes actually available to + // read. virtual size_t in_avail() const = 0; + // Return the number of bytes read by the last read operation. + size_t read_count() const { return m_read_count; } + bool has_error() const { return m_error; } void clear_error() { m_error = false; } @@ -69,16 +121,27 @@ class raw_istream { protected: void error_detected() { m_error = true; } + void set_read_count(size_t count) { m_read_count = count; } private: virtual void read_impl(void* data, size_t len) = 0; bool m_error = false; + size_t m_read_count = 0; }; class raw_mem_istream : public raw_istream { public: - explicit raw_mem_istream(StringRef mem); + // not const as we don't want to allow temporaries + explicit raw_mem_istream(std::string& str) + : raw_mem_istream(str.data(), str.size()) {} + explicit raw_mem_istream(ArrayRef mem) + : raw_mem_istream(mem.data(), mem.size()) {} + explicit raw_mem_istream(ArrayRef mem) + : raw_mem_istream(reinterpret_cast(mem.data()), mem.size()) { + } + explicit raw_mem_istream(const char* str) + : m_cur(str), m_left(std::strlen(str)) {} raw_mem_istream(const char* mem, size_t len) : m_cur(mem), m_left(len) {} void close() override; size_t in_avail() const override;