// Copyright (c) 2012 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 "update_engine/filesystem_copier_action.h"

#include <sys/stat.h>
#include <sys/types.h>
#include <errno.h>
#include <fcntl.h>

#include <algorithm>
#include <cstdlib>
#include <map>
#include <string>
#include <vector>

#include <base/string_util.h>
#include <base/stringprintf.h>
#include <gio/gio.h>
#include <gio/gunixinputstream.h>
#include <gio/gunixoutputstream.h>
#include <glib.h>

#include "update_engine/filesystem_iterator.h"
#include "update_engine/subprocess.h"
#include "update_engine/utils.h"

using std::map;
using std::min;
using std::string;
using std::vector;

namespace chromeos_update_engine {

namespace {
const off_t kCopyFileBufferSize = 128 * 1024;
}  // namespace {}

FilesystemCopierAction::FilesystemCopierAction(
    bool copying_kernel_install_path, bool verify_hash,
    unsigned max_rewrite_attempts)
    : copying_kernel_install_path_(copying_kernel_install_path),
      verify_hash_(verify_hash),
      src_stream_(NULL),
      dst_stream_(NULL),
      read_done_(false),
      failed_(false),
      cancelled_(false),
      filesystem_size_(kint64max),
      total_bytes_written_(0),
      num_curr_rewrite_attempts_(0),
      num_total_rewrite_attempts_(0),
      max_rewrite_attempts_(max_rewrite_attempts),
      is_debug_last_read_(false) {
  // A lot of code works on the implicit assumption that processing is done on
  // exactly 2 ping-pong buffers.
  COMPILE_ASSERT(arraysize(buffer_) == 2 &&
                 arraysize(buffer_state_) == 2 &&
                 arraysize(buffer_valid_size_) == 2 &&
                 arraysize(canceller_) == 2,
                 ping_pong_buffers_not_two);
  for (int i = 0; i < 2; ++i) {
    buffer_state_[i] = kBufferStateEmpty;
    buffer_valid_size_[i] = 0;
    canceller_[i] = NULL;
  }
}

void FilesystemCopierAction::PerformAction() {
  // Will tell the ActionProcessor we've failed if we return.
  ScopedActionCompleter abort_action_completer(processor_, this);

  if (!HasInputObject()) {
    LOG(ERROR) << "FilesystemCopierAction missing input object.";
    return;
  }
  install_plan_ = GetInputObject();
  if (!verify_hash_ && install_plan_.is_resume) {
    // No copy or hash verification needed. Done!
    if (HasOutputPipe())
      SetOutputObject(install_plan_);
    abort_action_completer.set_code(kActionCodeSuccess);
    return;
  }

  const string destination = copying_kernel_install_path_ ?
      install_plan_.kernel_install_path :
      install_plan_.install_path;
  string source = verify_hash_ ? destination : copy_source_;
  if (source.empty()) {
    source = copying_kernel_install_path_ ?
        utils::BootKernelDevice(utils::BootDevice()) :
        utils::BootDevice();
  }
  int src_fd = open(source.c_str(), O_RDONLY);
  if (src_fd < 0) {
    PLOG(ERROR) << "Unable to open " << source << " for reading:";
    return;
  }

  if (!verify_hash_) {
    int dst_fd = open(destination.c_str(),
                      O_WRONLY | O_TRUNC | O_CREAT,
                    0644);
    if (dst_fd < 0) {
      close(src_fd);
      PLOG(ERROR) << "Unable to open " << install_plan_.install_path
                  << " for writing:";
      return;
    }

    dst_stream_ = g_unix_output_stream_new(dst_fd, TRUE);
  }

  DetermineFilesystemSize(src_fd);
  src_stream_ = g_unix_input_stream_new(src_fd, TRUE);

  for (int i = 0; i < 2; i++) {
    buffer_[i].resize(kCopyFileBufferSize);
    canceller_[i] = g_cancellable_new();
  }

  // Start the first read.
  SpawnAsyncActions();

  abort_action_completer.set_should_complete(false);
}

void FilesystemCopierAction::TerminateProcessing() {
  for (int i = 0; i < 2; i++) {
    if (canceller_[i]) {
      g_cancellable_cancel(canceller_[i]);
    }
  }
}

void FilesystemCopierAction::Cleanup(ActionExitCode code) {
  for (int i = 0; i < 2; i++) {
    g_object_unref(canceller_[i]);
    canceller_[i] = NULL;
  }
  g_object_unref(src_stream_);
  src_stream_ = NULL;
  if (dst_stream_) {
    g_object_unref(dst_stream_);
    dst_stream_ = NULL;
  }
  if (cancelled_)
    return;
  if (code == kActionCodeSuccess && HasOutputPipe())
    SetOutputObject(install_plan_);
  processor_->ActionComplete(this, code);
}

void FilesystemCopierAction::AsyncReadReadyCallback(GObject *source_object,
                                                    GAsyncResult *res) {
  int index = buffer_state_[0] == kBufferStateReading ? 0 : 1;
  CHECK(buffer_state_[index] == kBufferStateReading);

  GError* error = NULL;
  CHECK(canceller_[index]);
  cancelled_ = g_cancellable_is_cancelled(canceller_[index]) == TRUE;

  ssize_t bytes_read = g_input_stream_read_finish(src_stream_, res, &error);

  // TODO(garnold) this is debugging code which needs to be removed once we
  // figure out the reasons for write I/O error.
  ssize_t bytes_expected = std::min(static_cast<int64_t>(buffer_[0].size()),
                                    filesystem_size_);
  if (!is_debug_last_read_ && bytes_read > 0 && bytes_read < bytes_expected) {
    LOG(INFO) << "[debug] read fewer bytes than expected ("
              << bytes_read << " < " << bytes_expected
              << "), this is likely the final read";
    is_debug_last_read_ = true;
  }

  if (bytes_read < 0) {
    LOG(ERROR) << "Read failed: " << utils::GetAndFreeGError(&error);
    failed_ = true;
    buffer_state_[index] = kBufferStateEmpty;
  } else if (bytes_read == 0) {
    read_done_ = true;
    buffer_state_[index] = kBufferStateEmpty;
  } else {
    buffer_valid_size_[index] = bytes_read;
    buffer_state_[index] = kBufferStateFull;
    filesystem_size_ -= bytes_read;
  }
  SpawnAsyncActions();

  if (bytes_read > 0) {
    // If read_done_ is set, SpawnAsyncActions may finalize the hash so the hash
    // update below would happen too late.
    CHECK(!read_done_);
    if (!hasher_.Update(buffer_[index].data(), bytes_read)) {
      LOG(ERROR) << "Unable to update the hash.";
      failed_ = true;
    }
    if (verify_hash_) {
      buffer_state_[index] = kBufferStateEmpty;
    }
  }
}

void FilesystemCopierAction::StaticAsyncReadReadyCallback(
    GObject *source_object,
    GAsyncResult *res,
    gpointer user_data) {
  reinterpret_cast<FilesystemCopierAction*>(user_data)->
      AsyncReadReadyCallback(source_object, res);
}

void FilesystemCopierAction::AsyncWriteReadyCallback(GObject *source_object,
                                                     GAsyncResult *res) {
  int index = buffer_state_[0] == kBufferStateWriting ? 0 : 1;
  CHECK(buffer_state_[index] == kBufferStateWriting);

  GError* error = NULL;
  CHECK(canceller_[index]);
  cancelled_ = g_cancellable_is_cancelled(canceller_[index]) == TRUE;

  ssize_t bytes_written = g_output_stream_write_finish(dst_stream_,
                                                       res,
                                                       &error);

  if (bytes_written < static_cast<ssize_t>(buffer_valid_size_[index])) {
    if (bytes_written < 0) {
      LOG(ERROR) << "Write error: " << utils::GetAndFreeGError(&error);
    } else {
      LOG(ERROR) << "Wrote too few bytes: " << bytes_written
                 << " instead of " << buffer_valid_size_[index];
    }

    bool do_fail = true;  // fail, unless we can retry

    if (cancelled_) {
      LOG(ERROR) << "write operation cancelled, copying failed";
    } else if (num_curr_rewrite_attempts_ < max_rewrite_attempts_) {
      // Try again: mark buffer as full, reposition the output stream.
      num_curr_rewrite_attempts_++;
      num_total_rewrite_attempts_++;
      off_t ret =
          lseek(g_unix_output_stream_get_fd(G_UNIX_OUTPUT_STREAM(dst_stream_)),
                static_cast<off_t>(total_bytes_written_),
                SEEK_SET);
      if (ret == static_cast<off_t>(total_bytes_written_)) {
        buffer_state_[index] = kBufferStateFull;
        do_fail = false;
        LOG(INFO) << "attempting to rewrite current buffer at offset "
                  << total_bytes_written_ << " (" << num_curr_rewrite_attempts_
                  << " of " << max_rewrite_attempts_ << " attempts)";
      } else {
        PLOG(ERROR)
            << "can't reposition output stream for rewrite, copying failed";
      }
    } else {
      LOG(ERROR) << "rewrite attempts exhausted (" << max_rewrite_attempts_
                 << "), copying failed";
    }

    if (do_fail)
      failed_ = true;
  } else {
    total_bytes_written_ += static_cast<size_t>(bytes_written);
    num_curr_rewrite_attempts_ = 0;
  }

  if (buffer_state_[index] == kBufferStateWriting)
    buffer_state_[index] = kBufferStateEmpty;

  SpawnAsyncActions();
}

void FilesystemCopierAction::StaticAsyncWriteReadyCallback(
    GObject *source_object,
    GAsyncResult *res,
    gpointer user_data) {
  reinterpret_cast<FilesystemCopierAction*>(user_data)->
      AsyncWriteReadyCallback(source_object, res);
}

void FilesystemCopierAction::SpawnAsyncActions() {
  bool reading = false;
  bool writing = false;
  for (int i = 0; i < 2; i++) {
    if (buffer_state_[i] == kBufferStateReading) {
      reading = true;
    }
    if (buffer_state_[i] == kBufferStateWriting) {
      writing = true;
    }
  }
  if (failed_ || cancelled_) {
    if (!reading && !writing) {
      Cleanup(kActionCodeError);
    }
    return;
  }
  for (int i = 0; i < 2; i++) {
    if (!reading && !read_done_ && buffer_state_[i] == kBufferStateEmpty) {
      int64_t bytes_to_read = std::min(static_cast<int64_t>(buffer_[0].size()),
                                       filesystem_size_);
      g_input_stream_read_async(
          src_stream_,
          buffer_[i].data(),
          bytes_to_read,
          G_PRIORITY_DEFAULT,
          canceller_[i],
          &FilesystemCopierAction::StaticAsyncReadReadyCallback,
          this);
      reading = true;
      buffer_state_[i] = kBufferStateReading;
    } else if (!writing && !verify_hash_ &&
               buffer_state_[i] == kBufferStateFull) {
      // TODO(garnold) debugging code, to be removed.
      if (is_debug_last_read_) {
        LOG(INFO) << "[debug] spawning async write of "
                  << buffer_valid_size_[i] << " bytes";
      }

      g_output_stream_write_async(
          dst_stream_,
          buffer_[i].data(),
          buffer_valid_size_[i],
          G_PRIORITY_DEFAULT,
          canceller_[i],
          &FilesystemCopierAction::StaticAsyncWriteReadyCallback,
          this);
      writing = true;
      buffer_state_[i] = kBufferStateWriting;
    }
  }
  if (!reading && !writing) {
    // We're done!
    ActionExitCode code = kActionCodeSuccess;
    // TODO(garnold) we're signaling a failure if buffer rewriting was attempted
    // during the copy action; this is necessary for keeping our unit tests
    // failing, allowing us to diagnose the root cause for write I/O errors.
    // When diagnosis is complete, this error return code needs to be reversed.
    if (num_total_rewrite_attempts_) {
      code = kActionCodeError;
      LOG(ERROR) << "[debug] rewrite was attempted "
                 << num_total_rewrite_attempts_ << " time(s), copying failed";
    }
    if (hasher_.Finalize()) {
      LOG(INFO) << "Hash: " << hasher_.hash();
      if (verify_hash_) {
        if (copying_kernel_install_path_) {
          if (install_plan_.kernel_hash != hasher_.raw_hash()) {
            code = kActionCodeNewKernelVerificationError;
            LOG(ERROR) << "New kernel verification failed.";
          }
        } else {
          if (install_plan_.rootfs_hash != hasher_.raw_hash()) {
            code = kActionCodeNewRootfsVerificationError;
            LOG(ERROR) << "New rootfs verification failed.";
          }
        }
      } else {
        if (copying_kernel_install_path_) {
          install_plan_.kernel_hash = hasher_.raw_hash();
        } else {
          install_plan_.rootfs_hash = hasher_.raw_hash();
        }
      }
    } else {
      LOG(ERROR) << "Unable to finalize the hash.";
      code = kActionCodeError;
    }
    Cleanup(code);
  }
}

void FilesystemCopierAction::DetermineFilesystemSize(int fd) {
  if (verify_hash_) {
    filesystem_size_ = copying_kernel_install_path_ ?
        install_plan_.kernel_size : install_plan_.rootfs_size;
    LOG(INFO) << "Filesystem size: " << filesystem_size_;
    return;
  }
  filesystem_size_ = kint64max;
  int block_count = 0, block_size = 0;
  if (!copying_kernel_install_path_ &&
      utils::GetFilesystemSizeFromFD(fd, &block_count, &block_size)) {
    filesystem_size_ = static_cast<int64_t>(block_count) * block_size;
    LOG(INFO) << "Filesystem size: " << filesystem_size_ << " bytes ("
              << block_count << "x" << block_size << ").";
  }
}

}  // namespace chromeos_update_engine