1// Copyright 2014 The Chromium Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5#include "chrome/utility/image_writer/image_writer.h"
6
7#include "base/memory/aligned_memory.h"
8#include "chrome/utility/image_writer/error_messages.h"
9#include "chrome/utility/image_writer/image_writer_handler.h"
10#include "content/public/utility/utility_thread.h"
11
12#if defined(OS_MACOSX)
13#include "chrome/utility/image_writer/disk_unmounter_mac.h"
14#endif
15
16namespace image_writer {
17
18// Since block devices like large sequential access and IPC is expensive we're
19// doing work in 1MB chunks.
20const int kBurningBlockSize = 1 << 20;  // 1 MB
21const int kMemoryAlignment = 4096;
22
23ImageWriter::ImageWriter(ImageWriterHandler* handler,
24                         const base::FilePath& image_path,
25                         const base::FilePath& device_path)
26    : image_path_(image_path),
27      device_path_(device_path),
28      bytes_processed_(0),
29      running_(false),
30      handler_(handler) {}
31
32ImageWriter::~ImageWriter() {
33#if defined(OS_WIN)
34  for (std::vector<HANDLE>::const_iterator it = volume_handles_.begin();
35       it != volume_handles_.end();
36       ++it) {
37    CloseHandle(*it);
38  }
39#endif
40}
41
42void ImageWriter::Write() {
43  if (!InitializeFiles()) {
44    return;
45  }
46
47  PostProgress(0);
48  PostTask(base::Bind(&ImageWriter::WriteChunk, AsWeakPtr()));
49}
50
51void ImageWriter::Verify() {
52  if (!InitializeFiles()) {
53    return;
54  }
55
56  PostProgress(0);
57  PostTask(base::Bind(&ImageWriter::VerifyChunk, AsWeakPtr()));
58}
59
60void ImageWriter::Cancel() {
61  running_ = false;
62  handler_->SendCancelled();
63}
64
65bool ImageWriter::IsRunning() const { return running_; }
66
67const base::FilePath& ImageWriter::GetImagePath() { return image_path_; }
68
69const base::FilePath& ImageWriter::GetDevicePath() { return device_path_; }
70
71void ImageWriter::PostTask(const base::Closure& task) {
72  base::MessageLoop::current()->PostTask(FROM_HERE, task);
73}
74
75void ImageWriter::PostProgress(int64 progress) {
76  handler_->SendProgress(progress);
77}
78
79void ImageWriter::Error(const std::string& message) {
80  running_ = false;
81  handler_->SendFailed(message);
82}
83
84bool ImageWriter::InitializeFiles() {
85  if (!image_file_.IsValid()) {
86    image_file_.Initialize(image_path_,
87                           base::File::FLAG_OPEN | base::File::FLAG_READ |
88                               base::File::FLAG_EXCLUSIVE_READ);
89
90    if (!image_file_.IsValid()) {
91      DLOG(ERROR) << "Unable to open file for read: " << image_path_.value();
92      Error(error::kOpenImage);
93      return false;
94    }
95  }
96
97  if (!device_file_.IsValid()) {
98    if (!OpenDevice()) {
99      Error(error::kOpenDevice);
100      return false;
101    }
102  }
103
104  bytes_processed_ = 0;
105  running_ = true;
106
107  return true;
108}
109
110void ImageWriter::WriteChunk() {
111  if (!IsRunning()) {
112    return;
113  }
114
115  // DASD buffers require memory alignment on some systems.
116  scoped_ptr<char, base::AlignedFreeDeleter> buffer(static_cast<char*>(
117      base::AlignedAlloc(kBurningBlockSize, kMemoryAlignment)));
118  memset(buffer.get(), 0, kBurningBlockSize);
119
120  int bytes_read = image_file_.Read(bytes_processed_, buffer.get(),
121                                    kBurningBlockSize);
122
123  if (bytes_read > 0) {
124    // Always attempt to write a whole block, as writing DASD requires sector-
125    // aligned writes to devices.
126    int bytes_to_write = bytes_read + (kMemoryAlignment - 1) -
127                         (bytes_read - 1) % kMemoryAlignment;
128    DCHECK_EQ(0, bytes_to_write % kMemoryAlignment);
129    int bytes_written =
130        device_file_.Write(bytes_processed_, buffer.get(), bytes_to_write);
131
132    if (bytes_written < bytes_read) {
133      Error(error::kWriteImage);
134      return;
135    }
136
137    bytes_processed_ += bytes_read;
138    PostProgress(bytes_processed_);
139
140    PostTask(base::Bind(&ImageWriter::WriteChunk, AsWeakPtr()));
141  } else if (bytes_read == 0) {
142    // End of file.
143    device_file_.Flush();
144    running_ = false;
145    handler_->SendSucceeded();
146  } else {
147    // Unable to read entire file.
148    Error(error::kReadImage);
149  }
150}
151
152void ImageWriter::VerifyChunk() {
153  if (!IsRunning()) {
154    return;
155  }
156
157  scoped_ptr<char[]> image_buffer(new char[kBurningBlockSize]);
158  // DASD buffers require memory alignment on some systems.
159  scoped_ptr<char, base::AlignedFreeDeleter> device_buffer(static_cast<char*>(
160      base::AlignedAlloc(kBurningBlockSize, kMemoryAlignment)));
161
162  int bytes_read = image_file_.Read(bytes_processed_, image_buffer.get(),
163                                    kBurningBlockSize);
164
165  if (bytes_read > 0) {
166    if (device_file_.Read(bytes_processed_,
167                          device_buffer.get(),
168                          kBurningBlockSize) < bytes_read) {
169      LOG(ERROR) << "Failed to read " << bytes_read << " bytes of "
170                 << "device at offset " << bytes_processed_;
171      Error(error::kReadDevice);
172      return;
173    }
174
175    if (memcmp(image_buffer.get(), device_buffer.get(), bytes_read) != 0) {
176      LOG(ERROR) << "Write verification failed when comparing " << bytes_read
177                 << " bytes at " << bytes_processed_;
178      Error(error::kVerificationFailed);
179      return;
180    }
181
182    bytes_processed_ += bytes_read;
183    PostProgress(bytes_processed_);
184
185    PostTask(base::Bind(&ImageWriter::VerifyChunk, AsWeakPtr()));
186  } else if (bytes_read == 0) {
187    // End of file.
188    handler_->SendSucceeded();
189    running_ = false;
190  } else {
191    // Unable to read entire file.
192    LOG(ERROR) << "Failed to read " << kBurningBlockSize << " bytes of image "
193               << "at offset " << bytes_processed_;
194    Error(error::kReadImage);
195  }
196}
197
198}  // namespace image_writer
199