[Paint Preview] Allow FileWStream size restriction

This CL simply augments FileWStream to support capping the file size.

This work is a precursor to allowing a limit on the total capture size.

Bug: 1071446
Change-Id: I38590bd52e50221f1007d92d673aca74f127158b
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2152768Reviewed-by: Mehran Mahmoudi <mahmoudi@chromium.org>
Commit-Queue: Calder Kitagawa <ckitagawa@chromium.org>
Cr-Commit-Position: refs/heads/master@{#760090}

components/paint_preview/common/file_stream.cc

@@ -9,25 +9,56 @@
 
 namespace paint_preview {
 
+namespace {
+
+bool ShouldWrite(size_t current_size, size_t max_size, size_t added_size) {
+  // If |current_size| + |added_size| overflow then don't write.
+  if (std::numeric_limits<size_t>::max() - current_size < added_size)
+    return false;
+  return max_size >= current_size + added_size;
+}
+
+}  // namespace
+
 // FileWStream
 
 FileWStream::FileWStream(base::File file)
-    : file_(std::move(file)), bytes_written_(0) {
+    : file_(std::move(file)),
+      max_size_(0),
+      bytes_written_(0),
+      has_write_failed_(false) {
+  DCHECK(file_.IsValid());
+}
+
+FileWStream::FileWStream(base::File file, size_t max_size)
+    : file_(std::move(file)),
+      max_size_(max_size),
+      bytes_written_(0),
+      has_write_failed_(false) {
   DCHECK(file_.IsValid());
 }
+
 // Close() is called in the destructor of |file_|.
 FileWStream::~FileWStream() = default;
 
 bool FileWStream::write(const void* buffer, size_t size) {
-  if (!file_.IsValid())
+  if (!file_.IsValid() || has_write_failed_)
+    return false;
+  if (max_size_ && !ShouldWrite(bytes_written_, max_size_, size)) {
+    has_write_failed_ = true;
     return false;
+  }
   int bytes =
       file_.WriteAtCurrentPos(reinterpret_cast<const char*>(buffer), size);
-  if (bytes < 0)
+  if (bytes < 0) {
+    has_write_failed_ = true;
     return false;
+  }
   bytes_written_ += bytes;
-  if (static_cast<size_t>(bytes) != size)
+  if (static_cast<size_t>(bytes) != size) {
+    has_write_failed_ = true;
     return false;
+  }
   return true;
 }
 

components/paint_preview/common/file_stream.h

@@ -14,7 +14,8 @@ namespace paint_preview {
 class FileWStream : public SkWStream {
  public:
   // Note: |file| must support writing.
-  FileWStream(base::File file);
+  explicit FileWStream(base::File file);
+  FileWStream(base::File file, size_t max_size);
   ~FileWStream() override;
 
   bool write(const void* buffer, size_t size) override;
@@ -24,9 +25,13 @@ class FileWStream : public SkWStream {
   // Closes the file (occurs automatically on destruction).
   void Close();
 
+  bool DidWriteFail() const { return has_write_failed_; }
+
  private:
   base::File file_;
+  size_t max_size_;
   size_t bytes_written_;
+  bool has_write_failed_;
 
   FileWStream(const FileWStream&) = delete;
   FileWStream& operator=(const FileWStream&) = delete;
@@ -37,7 +42,7 @@ class FileWStream : public SkWStream {
 class FileRStream : public SkStream {
  public:
   // Note: |file| must support reading. It *cannot* be modified while streaming.
-  FileRStream(base::File file);
+  explicit FileRStream(base::File file);
   ~FileRStream() override;
 
   size_t read(void* buffer, size_t size) override;

components/paint_preview/common/file_stream_unittest.cc

@@ -24,6 +24,7 @@ TEST(PaintPreviewFileStreamTest, TestWriteRead) {
   wstream.flush();
   wstream.Close();
   EXPECT_EQ(wstream.bytesWritten(), test_data.size());
+  EXPECT_FALSE(wstream.DidWriteFail());
   base::File read_file(file_path, base::File::FLAG_OPEN |
                                       base::File::FLAG_READ |
                                       base::File::FLAG_EXCLUSIVE_READ);
@@ -43,10 +44,59 @@ TEST(PaintPreviewFileStreamTest, TestWriteFail) {
   base::FilePath file_path = temp_dir.GetPath().AppendASCII("test_file");
   std::vector<uint8_t> test_data = {0, 1, 2, 3, 4, 5, 8, 9};
 
-  base::File write_file(file_path,
-                        base::File::FLAG_OPEN_ALWAYS | base::File::FLAG_READ);
-  FileWStream wstream(std::move(write_file));
+  base::File bad_write_file(
+      file_path, base::File::FLAG_OPEN_ALWAYS | base::File::FLAG_READ);
+  FileWStream wstream(std::move(bad_write_file));
+  EXPECT_FALSE(wstream.write(test_data.data(), test_data.size()));
+  EXPECT_TRUE(wstream.DidWriteFail());
+}
+
+// Test writing beyond max.
+TEST(PaintPreviewFileStreamTest, TestWriteFailCapped) {
+  base::ScopedTempDir temp_dir;
+  ASSERT_TRUE(temp_dir.CreateUniqueTempDir());
+  base::FilePath file_path = temp_dir.GetPath().AppendASCII("test_file");
+  std::vector<uint8_t> test_data = {0, 1, 2, 3, 4, 5, 8, 9};
+
+  base::File write_file(
+      file_path, base::File::FLAG_CREATE_ALWAYS | base::File::FLAG_WRITE);
+  FileWStream wstream(std::move(write_file), test_data.size());
+  EXPECT_TRUE(wstream.write(test_data.data(), test_data.size()));
+  EXPECT_FALSE(wstream.DidWriteFail());
+  EXPECT_FALSE(wstream.write(test_data.data(), test_data.size()));
+  EXPECT_TRUE(wstream.DidWriteFail());
+}
+
+// Test writing beyond max on first write.
+TEST(PaintPreviewFileStreamTest, TestWriteFailCappedFirstWrite) {
+  base::ScopedTempDir temp_dir;
+  ASSERT_TRUE(temp_dir.CreateUniqueTempDir());
+  base::FilePath file_path = temp_dir.GetPath().AppendASCII("test_file");
+  std::vector<uint8_t> test_data = {0, 1, 2, 3, 4, 5, 8, 9};
+
+  base::File write_file(
+      file_path, base::File::FLAG_CREATE_ALWAYS | base::File::FLAG_WRITE);
+  FileWStream wstream(std::move(write_file), 4);
   EXPECT_FALSE(wstream.write(test_data.data(), test_data.size()));
+  EXPECT_TRUE(wstream.DidWriteFail());
+}
+
+// Test writing beyond max (with overflow).
+TEST(PaintPreviewFileStreamTest, TestWriteFailOverflow) {
+  base::ScopedTempDir temp_dir;
+  ASSERT_TRUE(temp_dir.CreateUniqueTempDir());
+  base::FilePath file_path = temp_dir.GetPath().AppendASCII("test_file");
+  std::vector<uint8_t> test_data = {0, 1, 2, 3, 4, 5, 8, 9};
+
+  base::File write_file(
+      file_path, base::File::FLAG_CREATE_ALWAYS | base::File::FLAG_WRITE);
+  FileWStream wstream(std::move(write_file),
+                      std::numeric_limits<size_t>::max());
+  EXPECT_TRUE(wstream.write(test_data.data(), test_data.size()));
+  EXPECT_FALSE(wstream.DidWriteFail());
+  EXPECT_FALSE(
+      wstream.write(test_data.data(), std::numeric_limits<size_t>::max()));
+  EXPECT_TRUE(wstream.DidWriteFail());
 }
 
 // Test reading to skip.