src/diskimage_installer.cc

// Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "memento_common.h"
#include "diskimage_installer.h"

#define __STDC_FORMAT_MACROS // PRId64
#include <inttypes.h>
#include <stdlib.h>

#include <algorithm>
#include <string>
#include <sstream>

using std::string;
using std::stringstream;

namespace chromeos_memento_updater {

DiskImageInstaller::DiskImageInstaller(const string& running_dir,
                                       const string& board,
                                       const string& image_path,
                                       DiskImageInstallFormat format)
    : GeneralInstaller(running_dir, board) {
  image_path_ = image_path;
  format_ = format;
  allow_partition_options_ = true;
}

DiskImageInstaller::~DiskImageInstaller() {}

int DiskImageInstaller::DecompressImage(FILE** fp) {
  // TODO(chunyen): We may want to rewrite this with zlib, to reduce the
  // required disk space, and to support file not end with .zip.
  const string kZip = ".zip";
  size_t path_length = image_path_.length();
  // Return if file path does not end with .zip
  if (path_length <= kZip.length() ||
      image_path_.compare(path_length - kZip.length(),
                          kZip.length(), kZip) != 0) {
    return error("The path of file in zip format must ends with .zip");
  }
  // Unzip with gzip.
  if (system(("gzip -d -f -S .zip " + image_path_).c_str()) != 0) {
    return error("Fail to unzip image");
  }
  // Open the unzipped file (the .zip is removed by gzip).
  *fp = fopen(image_path_.substr(0, path_length - kZip.length()).c_str(), "r");
  format_ = kImage;
  image_path_ = image_path_.substr(0, path_length - kZip.length());
  return 0;
}

int DiskImageInstaller::PrepareImage(const int channel_id,
                                     FILE* out_pipe) {
  FILE* fp;
  switch (format_) {
    case kZip:
      if (DecompressImage(&fp) != 0) {
       return error("Fail to decompress image");
      };
      break;
    case kImage:
      fp = fopen(image_path_.c_str(), "rb");
      if (fp == NULL) {
        return error("Fail to read image");
      }
      break;
    default:
      return error("Unrecognizable format");
  }

  int64_t offset = GetImageOffsetInBlock(image_path_, channel_id) * kBlockSize;
  int64_t size = GetImageSizeInBlock(image_path_, channel_id) * kBlockSize;
  if (offset < 0 || size < 0) {
    return error("Fail to read cgpt header");
  }
  if (fseek(fp, offset, SEEK_SET) !=0 ) {
    fclose(fp);
    return error("Fail to fseek");
  }
  int64_t total_read = 0;
  int bytes_read;
  int bytes_write;
  const int64_t kBufSize = 4 * 1024 * 1024; // 4 MiB
  char* buffer = new char[kBufSize];
  while ((bytes_read =
          fread(buffer, 1 ,std::min(kBufSize, size - total_read), fp)) > 0) {
    total_read += bytes_read;
    bytes_write = 0;
    do {
      bytes_write += fwrite(&buffer[bytes_write], 1,
                            bytes_read - bytes_write, out_pipe);
    } while(!ferror(out_pipe) && bytes_write < bytes_read);
  }
  delete[] buffer;
  if (ferror(fp)) {
    fclose(fp);
    return error("Error reading image");
  }
  fclose(fp);
  if (ferror(out_pipe)) {
    return error("Error writing to pipe");
  }
  return 0;
}

// TODO(chunyen): Call vboot_reference/cgpt when it has a library.
int64_t DiskImageInstaller::GetImageOffsetInBlock(const string& image_path,
                                                  const int partition) {
  FILE* fp;
  int64_t offset;
  stringstream command;

  command << "cgpt show -b -i " << partition << " " << image_path;
  fp = popen(command.str().c_str(), "r");
  if (fp == NULL) {
    return -1;
  }
  if (fscanf(fp, "%"PRId64"", &offset) != 1) {
    pclose(fp);
    return -1;
  }
  pclose(fp);

  return offset;
}

int64_t DiskImageInstaller::GetImageSizeInBlock(const string& image_path,
                                                const int partition) {
  FILE* fp;
  int64_t size;
  stringstream command;

  command << "cgpt show -s -i " << partition << " " << image_path;
  fp = popen(command.str().c_str(), "r");
  if (fp == NULL) {
    return -1;
  }
  if (fscanf(fp, "%"PRId64"", &size) != 1) {
    pclose(fp);
    return -1;
  }
  pclose(fp);

  return size;
}

}  // namespace chromeos_memento_updater