Add FileWriter class to aos/util

This allows for a malloc-free writer for convenience.

Change-Id: I80cf6ac68f9190106fcb57d85cf758a5d491f04b
Signed-off-by: James Kuszmaul <james.kuszmaul@bluerivertech.com>
diff --git a/aos/util/BUILD b/aos/util/BUILD
index bd1dbab..3731965 100644
--- a/aos/util/BUILD
+++ b/aos/util/BUILD
@@ -336,6 +336,7 @@
         ":file",
         "//aos:realtime",
         "//aos/testing:googletest",
+        "//aos/testing:tmpdir",
     ],
 )
 
diff --git a/aos/util/file.cc b/aos/util/file.cc
index cdf0061..6c35627 100644
--- a/aos/util/file.cc
+++ b/aos/util/file.cc
@@ -36,20 +36,9 @@
 void WriteStringToFileOrDie(const std::string_view filename,
                             const std::string_view contents,
                             mode_t permissions) {
-  ScopedFD fd(open(::std::string(filename).c_str(),
-                   O_CREAT | O_WRONLY | O_TRUNC, permissions));
-  PCHECK(fd.get() != -1) << ": opening " << filename;
-  size_t size_written = 0;
-  while (size_written != contents.size()) {
-    const ssize_t result = write(fd.get(), contents.data() + size_written,
-                                 contents.size() - size_written);
-    PCHECK(result >= 0) << ": reading from " << filename;
-    if (result == 0) {
-      break;
-    }
-
-    size_written += result;
-  }
+  FileWriter writer(filename, permissions);
+  writer.WriteBytesOrDie(
+      {reinterpret_cast<const uint8_t *>(contents.data()), contents.size()});
 }
 
 bool MkdirPIfSpace(std::string_view path, mode_t mode) {
@@ -176,5 +165,47 @@
   return span;
 }
 
+FileWriter::FileWriter(std::string_view filename, mode_t permissions)
+    : file_(open(::std::string(filename).c_str(), O_WRONLY | O_CREAT | O_TRUNC,
+                 permissions)) {
+  PCHECK(file_.get() != -1) << ": opening " << filename;
+}
+
+FileWriter::WriteResult FileWriter::WriteBytes(
+    absl::Span<const uint8_t> bytes) {
+  size_t size_written = 0;
+  while (size_written != bytes.size()) {
+    const ssize_t result = write(file_.get(), bytes.data() + size_written,
+                                 bytes.size() - size_written);
+    if (result < 0) {
+      return {size_written, static_cast<int>(result)};
+    }
+    // Not really supposed to happen unless writing zero bytes without an error.
+    // See, e.g.,
+    // https://stackoverflow.com/questions/2176443/is-a-return-value-of-0-from-write2-in-c-an-error
+    if (result == 0) {
+      return {size_written, static_cast<int>(result)};
+    }
+
+    size_written += result;
+  }
+  return {size_written, static_cast<int>(size_written)};
+}
+
+FileWriter::WriteResult FileWriter::WriteBytes(std::string_view bytes) {
+  return WriteBytes(absl::Span<const uint8_t>{
+      reinterpret_cast<const uint8_t *>(bytes.data()), bytes.size()});
+}
+
+void FileWriter::WriteBytesOrDie(std::string_view bytes) {
+  WriteBytesOrDie(absl::Span<const uint8_t>{
+      reinterpret_cast<const uint8_t *>(bytes.data()), bytes.size()});
+}
+
+void FileWriter::WriteBytesOrDie(absl::Span<const uint8_t> bytes) {
+  PCHECK(bytes.size() == WriteBytes(bytes).bytes_written)
+      << ": Failed to write " << bytes.size() << " bytes.";
+}
+
 }  // namespace util
 }  // namespace aos
diff --git a/aos/util/file.h b/aos/util/file.h
index d2fb9fa..56479ce 100644
--- a/aos/util/file.h
+++ b/aos/util/file.h
@@ -80,6 +80,37 @@
   char buffer_[kBufferSize];
 };
 
+// Simple interface to allow opening a file for writing and then writing it
+// without any malloc's.
+// TODO(james): It may make sense to add a ReadBytes() interface here that can
+// take a memory buffer to fill, to avoid the templating required by the
+// self-managed buffer of FileReader<>.
+class FileWriter {
+ public:
+  // The result of an individual call to WriteBytes().
+  // Because WriteBytes() may repeatedly call write() when partial writes occur,
+  // it is possible for a non-zero number of bytes to have been written while
+  // still having an error because a late call to write() failed.
+  struct WriteResult {
+    // Total number of bytes successfully written to the file.
+    size_t bytes_written;
+    // If the write was successful (return_code > 0), equal to bytes_written.
+    // Otherwise, equal to the return value of the final call to write.
+    int return_code;
+  };
+
+  FileWriter(std::string_view filename, mode_t permissions = S_IRWXU);
+
+  WriteResult WriteBytes(absl::Span<const uint8_t> bytes);
+  WriteResult WriteBytes(std::string_view bytes);
+  void WriteBytesOrDie(absl::Span<const uint8_t> bytes);
+  void WriteBytesOrDie(std::string_view bytes);
+  int fd() const { return file_.get(); }
+
+ private:
+  aos::ScopedFD file_;
+};
+
 }  // namespace util
 }  // namespace aos
 
diff --git a/aos/util/file_test.cc b/aos/util/file_test.cc
index 8e76154..9b0def0 100644
--- a/aos/util/file_test.cc
+++ b/aos/util/file_test.cc
@@ -3,8 +3,9 @@
 #include <cstdlib>
 #include <string>
 
-#include "gtest/gtest.h"
 #include "aos/realtime.h"
+#include "aos/testing/tmpdir.h"
+#include "gtest/gtest.h"
 
 DECLARE_bool(die_on_malloc);
 
@@ -14,7 +15,7 @@
 
 // Basic test of reading a normal file.
 TEST(FileTest, ReadNormalFile) {
-  const ::std::string tmpdir(getenv("TEST_TMPDIR"));
+  const ::std::string tmpdir(aos::testing::TestTmpDir());
   const ::std::string test_file = tmpdir + "/test_file";
   ASSERT_EQ(0, system(("echo contents > " + test_file).c_str()));
   EXPECT_EQ("contents\n", ReadFileToStringOrDie(test_file));
@@ -30,7 +31,7 @@
 
 // Tests that the PathExists function works under normal conditions.
 TEST(FileTest, PathExistsTest) {
-  const std::string tmpdir(getenv("TEST_TMPDIR"));
+  const std::string tmpdir(aos::testing::TestTmpDir());
   const std::string test_file = tmpdir + "/test_file";
   // Make sure the test_file doesn't exist.
   unlink(test_file.c_str());
@@ -43,17 +44,65 @@
 
 // Basic test of reading a normal file.
 TEST(FileTest, ReadNormalFileNoMalloc) {
-  const ::std::string tmpdir(getenv("TEST_TMPDIR"));
+  const ::std::string tmpdir(aos::testing::TestTmpDir());
   const ::std::string test_file = tmpdir + "/test_file";
-  ASSERT_EQ(0, system(("echo 971 > " + test_file).c_str()));
+  // Make sure to include a string long enough to avoid small string
+  // optimization.
+  ASSERT_EQ(0, system(("echo 123456789 > " + test_file).c_str()));
 
   FileReader reader(test_file);
 
+  gflags::FlagSaver flag_saver;
   FLAGS_die_on_malloc = true;
   RegisterMallocHook();
   aos::ScopedRealtime realtime;
-  EXPECT_EQ("971\n", reader.ReadContents());
-  EXPECT_EQ(971, reader.ReadInt());
+  EXPECT_EQ("123456789\n", reader.ReadContents());
+  EXPECT_EQ(123456789, reader.ReadInt());
+}
+
+// Tests that we can write to a file without malloc'ing.
+TEST(FileTest, WriteNormalFileNoMalloc) {
+  const ::std::string tmpdir(aos::testing::TestTmpDir());
+  const ::std::string test_file = tmpdir + "/test_file";
+
+  FileWriter writer(test_file);
+
+  gflags::FlagSaver flag_saver;
+  FLAGS_die_on_malloc = true;
+  RegisterMallocHook();
+  FileWriter::WriteResult result;
+  {
+    aos::ScopedRealtime realtime;
+    result = writer.WriteBytes("123456789");
+  }
+  EXPECT_EQ(9, result.bytes_written);
+  EXPECT_EQ(9, result.return_code);
+  EXPECT_EQ("123456789", ReadFileToStringOrDie(test_file));
+}
+
+// Tests that if we fail to write a file that the error code propagates
+// correctly.
+TEST(FileTest, WriteFileError) {
+  const ::std::string tmpdir(aos::testing::TestTmpDir());
+  const ::std::string test_file = tmpdir + "/test_file";
+
+  // Open with only read permissions; this should cause things to fail.
+  FileWriter writer(test_file, S_IRUSR);
+
+  // Mess up the file management by closing the file descriptor.
+  PCHECK(0 == close(writer.fd()));
+
+  gflags::FlagSaver flag_saver;
+  FLAGS_die_on_malloc = true;
+  RegisterMallocHook();
+  FileWriter::WriteResult result;
+  {
+    aos::ScopedRealtime realtime;
+    result = writer.WriteBytes("123456789");
+  }
+  EXPECT_EQ(0, result.bytes_written);
+  EXPECT_EQ(-1, result.return_code);
+  EXPECT_EQ("", ReadFileToStringOrDie(test_file));
 }
 
 }  // namespace testing