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/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
  16 namespace image_writer {
  17
  18 // Since block devices like large sequential access and IPC is expensive we're
  19 // doing work in 1MB chunks.
  20 const int kBurningBlockSize = 1 << 20;  // 1 MB
  21 const int kMemoryAlignment = 4096;
  22
  23 ImageWriter::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
  32 ImageWriter::~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
  42 void ImageWriter::Write() {
  43   if (!InitializeFiles()) {
  44     return;
  45   }
  46
  47   PostProgress(0);
  48   PostTask(base::Bind(&ImageWriter::WriteChunk, AsWeakPtr()));
  49 }
  50
  51 void ImageWriter::Verify() {
  52   if (!InitializeFiles()) {
  53     return;
  54   }
  55
  56   PostProgress(0);
  57   PostTask(base::Bind(&ImageWriter::VerifyChunk, AsWeakPtr()));
  58 }
  59
  60 void ImageWriter::Cancel() {
  61   running_ = false;
  62   handler_->SendCancelled();
  63 }
  64
  65 bool ImageWriter::IsRunning() const { return running_; }
  66
  67 const base::FilePath& ImageWriter::GetImagePath() { return image_path_; }
  68
  69 const base::FilePath& ImageWriter::GetDevicePath() { return device_path_; }
  70
  71 void ImageWriter::PostTask(const base::Closure& task) {
  72   base::MessageLoop::current()->PostTask(FROM_HERE, task);
  73 }
  74
  75 void ImageWriter::PostProgress(int64 progress) {
  76   handler_->SendProgress(progress);
  77 }
  78
  79 void ImageWriter::Error(const std::string& message) {
  80   running_ = false;
  81   handler_->SendFailed(message);
  82 }
  83
  84 bool 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
 110 void 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     int bytes_written =
 129         device_file_.Write(bytes_processed_, buffer.get(), bytes_to_write);
 130
 131     if (bytes_written < bytes_read) {
 132       Error(error::kWriteImage);
 133       return;
 134     }
 135
 136     bytes_processed_ += bytes_read;
 137     PostProgress(bytes_processed_);
 138
 139     PostTask(base::Bind(&ImageWriter::WriteChunk, AsWeakPtr()));
 140   } else if (bytes_read == 0) {
 141     // End of file.
 142     device_file_.Flush();
 143     running_ = false;
 144     handler_->SendSucceeded();
 145   } else {
 146     // Unable to read entire file.
 147     Error(error::kReadImage);
 148   }
 149 }
 150
 151 void ImageWriter::VerifyChunk() {
 152   if (!IsRunning()) {
 153     return;
 154   }
 155
 156   scoped_ptr<char[]> image_buffer(new char[kBurningBlockSize]);
 157   // DASD buffers require memory alignment on some systems.
 158   scoped_ptr<char, base::AlignedFreeDeleter> device_buffer(static_cast<char*>(
 159       base::AlignedAlloc(kBurningBlockSize, kMemoryAlignment)));
 160
 161   int bytes_read = image_file_.Read(bytes_processed_, image_buffer.get(),
 162                                     kBurningBlockSize);
 163
 164   if (bytes_read > 0) {
 165     if (device_file_.Read(bytes_processed_,
 166                           device_buffer.get(),
 167                           kBurningBlockSize) < bytes_read) {
 168       LOG(ERROR) << "Failed to read " << bytes_read << " bytes of "
 169                  << "device at offset " << bytes_processed_;
 170       Error(error::kReadDevice);
 171       return;
 172     }
 173
 174     if (memcmp(image_buffer.get(), device_buffer.get(), bytes_read) != 0) {
 175       LOG(ERROR) << "Write verification failed when comparing " << bytes_read
 176                  << " bytes at " << bytes_processed_;
 177       Error(error::kVerificationFailed);
 178       return;
 179     }
 180
 181     bytes_processed_ += bytes_read;
 182     PostProgress(bytes_processed_);
 183
 184     PostTask(base::Bind(&ImageWriter::VerifyChunk, AsWeakPtr()));
 185   } else if (bytes_read == 0) {
 186     // End of file.
 187     handler_->SendSucceeded();
 188     running_ = false;
 189   } else {
 190     // Unable to read entire file.
 191     LOG(ERROR) << "Failed to read " << kBurningBlockSize << " bytes of image "
 192                << "at offset " << bytes_processed_;
 193     Error(error::kReadImage);
 194   }
 195 }
 196
 197 }  // namespace image_writer