chrome/utility/image_writer/image_writer.cc

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