--- /dev/null
+/*
+ * Copyright 2020 Google Inc.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met: redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer;
+ * redistributions in binary form must reproduce the above copyright
+ * notice, this list of conditions and the following disclaimer in the
+ * documentation and/or other materials provided with the distribution;
+ * neither the name of the copyright holders nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include <sys/mman.h>
+#include <sys/stat.h>
+#include <unistd.h>
+
+#include <cstdlib>
+#include <cstring>
+#include <sstream>
+#include <string>
+
+#include "args.hh"
+#include "command.hh"
+#include "dispatch_table.hh"
+
+uint64_t test_total_written;
+std::string test_host_file_name;
+
+std::vector<uint8_t> test_written_data;
+uint64_t test_max_buf_size;
+
+uint64_t
+test_m5_write_file(void *buffer, uint64_t len, uint64_t offset,
+ const char *filename)
+{
+ if (test_max_buf_size && len > test_max_buf_size)
+ len = test_max_buf_size;
+
+ test_total_written += len;
+
+ if (test_host_file_name == "")
+ test_host_file_name = filename;
+ else
+ EXPECT_EQ(test_host_file_name, filename);
+
+ if (offset == 0)
+ test_written_data.clear();
+
+ size_t required_size = offset + len;
+ if (test_written_data.size() < required_size)
+ test_written_data.resize(required_size);
+
+ memcpy(test_written_data.data() + offset, buffer, len);
+
+ return len;
+}
+
+DispatchTable dt = { .m5_write_file = &test_m5_write_file };
+
+std::string cout_output;
+
+bool
+run(std::initializer_list<std::string> arg_args)
+{
+ test_total_written = 0;
+ test_host_file_name = "";
+ test_written_data.clear();
+
+ Args args(arg_args);
+
+ // Redirect cout into a stringstream.
+ std::stringstream buffer;
+ std::streambuf *orig = std::cout.rdbuf(buffer.rdbuf());
+
+ bool res = Command::run(dt, args);
+
+ // Capture the contents of the stringstream and restore cout.
+ cout_output = buffer.str();
+ std::cout.rdbuf(orig);
+
+ return res;
+}
+
+class TempFile
+{
+ private:
+ size_t _size;
+ int fd;
+ std::string _path;
+ void *_buf;
+
+ public:
+ TempFile(size_t _size) : _size(_size)
+ {
+ // Generate a temporary filename.
+ char *tmp_name = strdup("/tmp/writefile.test.XXXXXXXX");
+ fd = mkstemp(tmp_name);
+ _path = tmp_name;
+ free(tmp_name);
+
+ // Make the file the appropriate length.
+ assert(!ftruncate(fd, _size));
+
+ // mmap the file.
+ _buf = mmap(nullptr, _size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
+ assert(_buf);
+
+ // Fill it with an incrementing 32 bit integers.
+
+ int chunk_size = sizeof(uint32_t);
+ size_t num_chunks = _size / chunk_size;
+ int leftovers = _size % chunk_size;
+
+ uint32_t *buf32 = (uint32_t *)_buf;
+ uint32_t val = 0;
+ for (size_t i = 0; i < num_chunks; i++)
+ *buf32++ = val++;
+ if (leftovers)
+ memcpy(buf32, &val, leftovers);
+
+ // Make sure our new contents are out there.
+ msync(_buf, _size, MS_SYNC | MS_INVALIDATE);
+ };
+
+ ~TempFile()
+ {
+ unlink(path().c_str());
+ close(fd);
+ }
+
+ size_t size() const { return _size; }
+ const std::string &path() const { return _path; }
+ const void *buf() const { return _buf; }
+
+ void
+ verify()
+ {
+ verify(path());
+ }
+
+ void
+ verify(const std::string &expected_path)
+ {
+ EXPECT_EQ(test_written_data.size(), size());
+ EXPECT_EQ(memcmp(test_written_data.data(), buf(), size()), 0);
+ EXPECT_EQ(test_host_file_name, expected_path);
+
+ std::ostringstream os;
+ os << "Opening \"" << path() << "\".";
+ EXPECT_THAT(cout_output, ::testing::HasSubstr(os.str()));
+ os.str("");
+ os << "Wrote " << size() << " bytes.";
+ EXPECT_THAT(cout_output, ::testing::HasSubstr(os.str()));
+ }
+};
+
+TEST(Writefile, NoArguments)
+{
+ // Call with no arguments.
+ EXPECT_FALSE(run({"writefile"}));
+ EXPECT_EQ(test_total_written, 0);
+}
+
+TEST(Writefile, ThreeArguments)
+{
+ // Call with no arguments.
+ EXPECT_FALSE(run({"writefile", "1", "2", "3"}));
+ EXPECT_EQ(test_total_written, 0);
+}
+
+TEST(Writefile, SmallFile)
+{
+ // Write a small file.
+ TempFile tmp(16);
+ test_max_buf_size = 0;
+ EXPECT_TRUE(run({"writefile", tmp.path()}));
+ tmp.verify();
+}
+
+TEST(Writefile, SmallFileHostName)
+{
+ // Write a small file with a different host file name.
+ TempFile tmp(16);
+ test_max_buf_size = 0;
+ std::string host_path = "/different/host/path";
+ EXPECT_TRUE(run({"writefile", tmp.path(), host_path}));
+ tmp.verify(host_path);
+}
+
+TEST(Writefile, MultipleChunks)
+{
+ // Write a file which will need to be split into multiple whole chunks.
+ TempFile tmp(256 * 1024 * 4);
+ test_max_buf_size = 0;
+ EXPECT_TRUE(run({"writefile", tmp.path()}));
+ tmp.verify();
+}
+
+TEST(Writefile, MultipleAndPartialChunks)
+{
+ // Write a file which will be split into some whole and one partial chunk.
+ TempFile tmp(256 * 1024 * 2 + 256);
+ test_max_buf_size = 0;
+ EXPECT_TRUE(run({"writefile", tmp.path()}));
+ tmp.verify();
+}
+
+TEST(Writefile, OddSizedChunks)
+{
+ // Write a file in chunks that aren't nicely aligned.
+ TempFile tmp(256 * 1024);
+ test_max_buf_size = 13;
+ EXPECT_TRUE(run({"writefile", tmp.path()}));
+ tmp.verify();
+}
+
+TEST(Writefile, CappedWriteSize)
+{
+ // Write a file, accepting less than the requested amount of data.
+ TempFile tmp(256 * 1024 * 2 + 256);
+ test_max_buf_size = 256;
+ EXPECT_TRUE(run({"writefile", tmp.path()}));
+ tmp.verify();
+}
+
+TEST(WritefileDeathTest, BadFile)
+{
+ EXPECT_EXIT(run({"writefile", "this is not a valid path#$#$://\\\\"}),
+ ::testing::ExitedWithCode(2), "Error opening ");
+}