Ensure that LoadSequentialFile() actually read the whole file

include/xgboost/json_io.h

@@ -9,12 +9,13 @@
 #include <vector>
 #include <memory>
 #include <string>
-#include <cinttypes>
 #include <utility>
 #include <map>
 #include <limits>
 #include <sstream>
 #include <locale>
+#include <cinttypes>
+#include <cstdio>
 
 namespace xgboost {
 /*
@@ -84,8 +85,10 @@ class JsonReader {
     std::string msg = "Expecting: \"";
     msg += c;
     msg += "\", got: \"";
-    if (got == -1) {
+    if (got == EOF) {
       msg += "EOF\"";
+    } else if (got == 0) {
+      msg += "\\0\"";
     } else {
       msg += std::to_string(got) + " \"";
     }

src/common/io.cc

@@ -7,9 +7,10 @@
 #include <unistd.h>
 #endif  // defined(__unix__)
 #include <algorithm>
-#include <cstdio>
+#include <fstream>
 #include <string>
 #include <utility>
+#include <cstdio>
 
 #include "xgboost/logging.h"
 #include "io.h"
@@ -120,27 +121,22 @@ std::string LoadSequentialFile(std::string fname) {
 #if defined(__linux__)
   posix_fadvise(fd, 0, 0, POSIX_FADV_SEQUENTIAL);
 #endif  // defined(__linux__)
-  ssize_t bytes_read = read(fd, &buffer[0], f_size_bytes);
-  if (bytes_read < 0) {
-    close(fd);
-    ReadErr();
+  ssize_t bytes_read = 0;
+  while (bytes_read < f_size_bytes) {
+    ssize_t result = read(fd, &buffer[bytes_read], f_size_bytes);
+    if (result < 0) {
+      close(fd);
+      ReadErr();
+    }
+    bytes_read += result;
   }
   close(fd);
+  buffer.back() = '\0';
 #else  // defined(__unix__)
-  FILE *f = fopen(fname.c_str(), "r");
-  if (f == NULL) {
-    std::string msg;
-    OpenErr();
-  }
-  fseek(f, 0, SEEK_END);
-  auto fsize = ftell(f);
-  fseek(f, 0, SEEK_SET);
-
-  buffer.resize(fsize + 1);
-  fread(&buffer[0], 1, fsize, f);
-  fclose(f);
+  std::ifstream ifs(fname);
+  buffer = std::string( (std::istreambuf_iterator<char>(ifs) ),
+                        (std::istreambuf_iterator<char>()) );
 #endif  // defined(__unix__)
-  buffer.back() = '\0';
   return buffer;
 }
 

src/common/json.cc

@@ -427,6 +427,8 @@ void JsonReader::Error(std::string msg) const {
   for (auto c : raw_portion) {
     if (c == '\n') {
       portion += "\\n";
+    } else if (c == '\0') {
+      portion += "\\0";
     } else {
       portion += c;
     }

tests/cpp/common/test_io.cc

@@ -2,6 +2,11 @@
  * Copyright (c) by XGBoost Contributors 2019
  */
 #include <gtest/gtest.h>
+#include <dmlc/filesystem.h>
+#include <atomic>
+#include <type_traits>
+#include <fstream>
+#include <cstdint>
 #include "../../../src/common/io.h"
 
 namespace xgboost {
@@ -39,5 +44,51 @@ TEST(IO, FixedSizeStream) {
     ASSERT_EQ(huge_buffer, out_buffer);
   }
 }
+
+
+#if SIZE_MAX == 0xFFFFFFFFFFFFFFFF  // Only run this test on 64-bit system
+TEST(IO, LoadSequentialFile) {
+  // bytes_read = 2147479552, f_size_bytes = 2896075944
+  const size_t nbyte = static_cast<size_t>(2896075944LL);
+  static_assert(sizeof(size_t) == 8, "Assumption failed: size_t was assumed to be 8-bytes long");
+  static_assert(std::is_same<size_t, std::string::size_type>::value,
+                "Assumption failed: size_type of std::string was assumed to be 8-bytes long");
+
+  dmlc::TemporaryDirectory tempdir;
+  std::string path = "/dev/shm/xgboost_test_io_big_file.txt";
+  {
+    std::ofstream f(path);
+    if (!f) {  // /dev/shm not present
+      LOG(INFO) << "No /dev/shm; using dmlc::TemporaryDirectory instead";
+      path = tempdir.path + "/xgboost_test_io_big_file.txt";
+      f = std::ofstream(path);
+    }
+    CHECK(f);
+    std::string str(nbyte, 'a');
+    CHECK_EQ(str.size(), nbyte);
+    f << str;
+  }
+  {
+    std::string str = LoadSequentialFile(path);
+    CHECK_EQ(str.size(), nbyte);
+    dmlc::OMPException omp_exc;
+    std::atomic<bool> success{true};
+    #pragma omp parallel for schedule(static)
+    for (int64_t i = 0; i < static_cast<int64_t>(nbyte); ++i) {
+      omp_exc.Run([&] {
+        if (str[i] != 'a' && success.load(std::memory_order_acquire)) {
+          success.store(false, std::memory_order_release);
+          LOG(FATAL) << "Big file got corrupted. Expected: str[" << i << "] = 'a', "
+            << "Actual: str[" << i << "] = '"
+            << (str[i] ? std::string(1, str[i]) : std::string("\\0")) << "'";
+        }
+      });
+    }
+    omp_exc.Rethrow();
+    CHECK(success.load(std::memory_order_acquire));
+  }
+}
+#endif  // SIZE_MAX == 0xFFFFFFFFFFFFFFFF
+
 }  // namespace common
 }  // namespace xgboost