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     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
 152 void 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