1 files changed, 50 insertions, 41 deletions

diff --git a/src/buffered_file.cpp b/src/buffered_file.cpp
index f5eb863..8d47eed 100644
--- a/src/buffered_file.cpp
+++ b/src/buffered_file.cpp
@@ -27,53 +27,71 @@
 #include "buffered_file.h"
 #include "exception.h"
 
+#include <algorithm>
+#include <cassert>
 #include <cstring>
 #include <filesystem>
 #include <fstream>
 
-BufferedFileReader::BufferedFileReader(std::filesystem::path path,
+BufferedFileReader::BufferedFileReader(std::istream &istream,
                                        size_t buffer_capacity)
-    : m_buffer_offset(0), m_buffer_capacity(buffer_capacity)
+    : m_istream(istream), m_buffer_offset(0), m_buffer_size(0),
+      m_buffer_capacity(buffer_capacity)
 {
-    m_file.open(path, std::ifstream::in | std::ifstream::binary);
-    if (!m_file) {
-        throw BufferedFileError("cannot open input file");
-    }
-
     try {
         m_buffer = std::make_unique<char[]>(m_buffer_capacity);
     } catch (const std::bad_alloc &e) {
         throw BufferedFileError("cannot allocate buffer for input file data");
     }
-
-    refill_buffer();
 };
 
 size_t
 BufferedFileReader::read(char *data, size_t data_size)
 {
-    const size_t size_left = m_buffer_size - m_buffer_offset;
-    if (data_size <= size_left) {
-        return read_buffer(data, data_size);
-    } else {
-        const size_t size_outside_buffer = data_size - size_left;
-        read_buffer(data, size_left);
-        const size_t read_outside =
-            read_file(data + size_left, size_outside_buffer);
+    size_t retry_count{0};
+    size_t offset{0};
+
+    while ((data_size > 0) && (retry_count < 2)) {
+        char *d;
+        const size_t r{tryRead(data_size, &d)};
+        if (r == 0) {
+            ++retry_count;
+            continue;
+        }
+
+        memcpy(data + offset, d, r);
+        offset += r;
+        data_size -= r;
+    }
+
+    return offset;
+}
+
+size_t
+BufferedFileReader::tryRead(size_t data_size, char **return_data)
+{
+    const size_t size_left{m_buffer_size - m_buffer_offset};
+    if (size_left == 0) {
         refill_buffer();
-        return size_left + read_outside;
+        if (m_buffer_size == 0) {
+            return 0;
+        }
     }
+    // There is at least one byte in the buffer
+    const size_t size_read{read_buffer(data_size, return_data)};
+    assert(size_read > 0);
+    return size_read;
 };
 
 size_t
-BufferedFileReader::read_buffer(char *data, size_t data_size)
+BufferedFileReader::read_buffer(size_t data_size, char **return_data)
 {
-    // Assumes that the caller makes sure there is enough data in the buffer
-    // to read
-    memcpy(data, reinterpret_cast<char *>(m_buffer.get()) + m_buffer_offset,
-           data_size);
-    m_buffer_offset += data_size;
-    return data_size;
+    *return_data = static_cast<char *>(m_buffer.get()) + m_buffer_offset;
+
+    const size_t size_left{m_buffer_size - m_buffer_offset};
+    const size_t size_read{std::min(data_size, size_left)};
+    m_buffer_offset += size_read;
+    return size_read;
 };
 
 void
@@ -86,28 +104,19 @@ BufferedFileReader::refill_buffer()
 size_t
 BufferedFileReader::read_file(char *data, size_t data_size)
 {
-    m_file.read(data, data_size);
+    m_istream.read(data, data_size);
 
-    if (!m_file.good() && !m_file.eof()) {
+    if (!m_istream.good() && !m_istream.eof()) {
         throw BufferedFileError("cannot read from file");
     }
 
-    return m_file.gcount();
+    return m_istream.gcount();
 };
 
-BufferedFileWriter::BufferedFileWriter(std::filesystem::path path,
+BufferedFileWriter::BufferedFileWriter(std::ostream &ostream,
                                        size_t buffer_capacity)
-    : m_buffer_size(0), m_buffer_capacity(buffer_capacity)
+    : m_ostream(ostream), m_buffer_size(0), m_buffer_capacity(buffer_capacity)
 {
-    /* When backing up, the output file is truncated to hold the
-     * new data
-     */
-    m_file.open(path, std::ifstream::out | std::ifstream::trunc |
-                          std::ifstream::binary);
-    if (!m_file) {
-        throw BufferedFileError("cannot open output file");
-    }
-
     try {
         m_buffer = std::make_unique<char[]>(m_buffer_capacity);
     } catch (const std::bad_alloc &e) {
@@ -154,8 +163,8 @@ BufferedFileWriter::flush_buffer()
 void
 BufferedFileWriter::write_file(const char *data, size_t data_size)
 {
-    m_file.write(data, data_size);
-    if (!m_file) {
+    m_ostream.write(data, data_size);
+    if (!m_ostream) {
         throw BufferedFileError("cannot write to output file");
     }
 };