// Copyright 2016 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 "crash-reporter/core-collector/coredump_writer.h"

#include <errno.h>
#include <fcntl.h>
#include <sys/statvfs.h>
#include <sys/types.h>
#include <sysexits.h>
#include <unistd.h>

#include <algorithm>
#include <cstdint>
#include <cstring>
#include <string>

#include <common/linux/elf_core_dump.h>

#include "crash-reporter/core-collector/logging.h"

using google_breakpad::ElfCoreDump;

namespace {

const size_t kMaxAbsCoredumpSize = 256 * 1024 * 1024;
const double kMaxRelCoredumpSize = 0.05;  // Percentage of free disk space.

struct ScopedFd {
  explicit ScopedFd(int fd) : fd(fd) {}
  ~ScopedFd() { close(fd); }

  bool IsValid() const { return fd >= 0; }
  operator int() const { return fd; }

  const int fd;
};

template <typename DataPtr>
using FdFunctionPtr = ssize_t (*)(int, DataPtr, size_t);

// Reads or writes exactly |count| bytes of |data| using |fd| as the source or
// destination, respectively.
template <typename DataPtr, typename BytePtr, FdFunctionPtr<DataPtr> op>
struct FdOperation {
  bool operator ()(int fd, DataPtr data, size_t count) {
    BytePtr ptr = reinterpret_cast<BytePtr>(data);
    while (count > 0) {
      const ssize_t n = TEMP_FAILURE_RETRY(op(fd, ptr, count));
      if (n < 0)
        return false;
      if (n == 0) {
        errno = EIO;
        return false;
      }
      ptr += n;
      count -= n;
    }
    return true;
  }
};

FdOperation<void *, uint8_t *, &read> ReadAllBlocking;
FdOperation<const void *, const uint8_t *, &write> WriteAllBlocking;

inline bool Seek(int fd, off_t offset) {
  return lseek(fd, offset, SEEK_SET) == offset;
}

inline int64_t GetFreeDiskSpace(const char *path) {
  struct statvfs stats;
  return TEMP_FAILURE_RETRY(statvfs(path, &stats)) != 0 ? -1 :
      static_cast<int64_t>(stats.f_bavail) * stats.f_frsize;
}

}  // namespace

class CoredumpWriter::Reader {
 public:
  explicit Reader(int fd) : fd_(fd) {}

  bool Read(void *buf, size_t count) {
    if (!ReadAllBlocking(fd_, buf, count))
      return false;
    bytes_read_ += count;
    return true;
  }

  bool CopyTo(int dest_fd, size_t count) {
    static const size_t kBufSize = 32 * 1024;
    char buf[kBufSize];
    while (count > 0) {
      const ssize_t n = TEMP_FAILURE_RETRY(
          read(fd_, buf, std::min(kBufSize, count)));
      if (n < 0)
        return false;
      if (n == 0) {
        errno = 0;
        break;
      }
      if (dest_fd >= 0 && !WriteAllBlocking(dest_fd, buf, n))
        return false;
      count -= n;
      bytes_read_ += n;
    }
    return count == 0;
  }

  bool Seek(size_t offset) {
    if (offset < bytes_read_)  // Cannot move backward.
      return false;
    return CopyTo(-1, offset - bytes_read_);
  }

 private:
  const int fd_;
  size_t bytes_read_ = 0;

  DISALLOW_COPY_AND_ASSIGN(Reader);
};

CoredumpWriter::CoredumpWriter(int fd,
                               const char *coredump_path,
                               const char *container_dir)
    : fd_(fd),
      coredump_path_(coredump_path),
      container_dir_(container_dir) {
}

int CoredumpWriter::WriteCoredump() {
  const ScopedFd dest(TEMP_FAILURE_RETRY(
      open(coredump_path_, O_WRONLY | O_CREAT | O_EXCL, S_IRUSR | S_IWUSR)));

  if (!dest.IsValid()) {
    PLOG_ERROR << "Failed to open '" << coredump_path_ << "'";
    return EX_CANTCREAT;
  }

  // Input core dump is generated by kernel's fs/binfmt_elf.c using this format:
  //
  //   ELF Header
  //   Program Header 1
  //   Program Header 2
  //   ...
  //   Program Header n
  //   Segment 1 (This segment's type should be PT_NOTE)
  //   Segment 2
  //   ...
  //   Segment n

  // Read ELF header, all program headers, and the first segment whose type is
  // PT_NOTE.
  Reader reader(fd_);
  Ehdr elf_header;
  std::vector<Phdr> program_headers;
  std::vector<char> note_buf;
  int error = ReadUntilNote(&reader, &elf_header, &program_headers, &note_buf);
  if (error != EX_OK)
    return error;

  // Get a set of address ranges occupied by mapped files from PT_NOTE segment.
  FileMappings file_mappings;
  if (!GetFileMappings(note_buf, &file_mappings))
    return EX_OSFILE;

  // Strip segments backed by mapped files, since they are not needed to
  // generate a minidump.
  std::vector<Phdr> stripped_program_headers;
  StripSegments(program_headers, file_mappings, &stripped_program_headers);

  // Calculate the core dump size limit.
  const int64_t free_disk_space = GetFreeDiskSpace(coredump_path_);
  if (free_disk_space < 0) {
    PLOG_ERROR << "Failed to get free disk space";
    return EX_OSERR;
  }
  const auto coredump_size_limit = std::min(kMaxAbsCoredumpSize,
      static_cast<size_t>(free_disk_space * kMaxRelCoredumpSize));

  // Calculate the output file size.
  const auto &last = stripped_program_headers.back();
  const auto expected_coredump_size = last.p_offset + last.p_filesz;
  if (expected_coredump_size > coredump_size_limit) {
    LOG_ERROR << "Exceeded maximum core dump size by "
              << expected_coredump_size - coredump_size_limit << " bytes";
    return EX_CANTCREAT;
  }

  // Write /proc files.
  error = WriteAuxv(note_buf);
  if (error != EX_OK) {
    LOG_ERROR << "Failed to write auxv";
    return error;
  }
  error = WriteMaps(program_headers, file_mappings);
  if (error != EX_OK) {
    LOG_ERROR << "Failed to write maps";
    return error;
  }

  // Write ELF header.
  if (!WriteAllBlocking(dest, &elf_header, sizeof(elf_header))) {
    PLOG_ERROR << "Failed to write ELF header";
    return EX_IOERR;
  }

  // Write program headers.
  for (size_t i = 0; i < stripped_program_headers.size(); ++i) {
    const Phdr &program_header = stripped_program_headers[i];
    const auto offset = sizeof(elf_header) + i * elf_header.e_phentsize;
    if (!Seek(dest, offset) ||
        !WriteAllBlocking(dest, &program_header, sizeof(program_header))) {
      PLOG_ERROR << "Failed to write program header";
      return EX_IOERR;
    }
  }

  // Write PT_NOTE segment.
  if (!Seek(dest, stripped_program_headers[0].p_offset) ||
      !WriteAllBlocking(dest, note_buf.data(), note_buf.size())) {
    PLOG_ERROR << "Failed to write PT_NOTE segment";
    return EX_IOERR;
  }

  // Write segments that were not stripped.
  for (size_t i = 1; i < stripped_program_headers.size(); ++i) {
    const Phdr &program_header = stripped_program_headers[i];
    if (program_header.p_filesz == 0)
      continue;
    const Phdr &program_header_original = program_headers[i];
    if (!reader.Seek(program_header_original.p_offset)) {
      PLOG_ERROR << "Failed to seek segment";
      return EX_IOERR;
    }
    if (!Seek(dest, program_header.p_offset) ||
        !reader.CopyTo(dest, program_header.p_filesz)) {
      PLOG_ERROR << "Failed to write segment";
      return EX_IOERR;
    }
  }

  return EX_OK;
}

int CoredumpWriter::ReadUntilNote(Reader *reader,
                                  Ehdr *elf_header,
                                  std::vector<Phdr> *program_headers,
                                  std::vector<char> *note_buf) {
  // Read ELF header.
  if (!reader->Read(elf_header, sizeof(*elf_header))) {
    PLOG_ERROR << "Failed to read ELF header";
    return EX_IOERR;
  }
  if (memcmp(elf_header->e_ident, ELFMAG, SELFMAG) != 0 ||
      elf_header->e_ident[EI_CLASS] != ElfCoreDump::kClass ||
      elf_header->e_version != EV_CURRENT ||
      elf_header->e_type != ET_CORE ||
      elf_header->e_ehsize != sizeof(Ehdr) ||
      elf_header->e_phentsize != sizeof(Phdr)) {
    LOG_ERROR << "Invalid ELF header";
    return EX_OSFILE;
  }

  // Read program headers.
  program_headers->resize(elf_header->e_phnum);
  if (!reader->Seek(elf_header->e_phoff) ||
      !reader->Read(program_headers->data(),
                    sizeof(Phdr) * program_headers->size())) {
    PLOG_ERROR << "Failed to read program headers";
    return EX_IOERR;
  }

  // The first segment should have type PT_NOTE. This assumption hinges on the
  // kernel's fs/binfmt_elf.c implementation.
  if (program_headers->empty() || (*program_headers)[0].p_type != PT_NOTE) {
    LOG_ERROR << "Failed to locate PT_NOTE segment";
    return EX_OSFILE;
  }
  const Phdr &note_program_header = (*program_headers)[0];

  // Read PT_NOTE segment.
  note_buf->resize(note_program_header.p_filesz);
  if (!reader->Seek(note_program_header.p_offset) ||
      !reader->Read(note_buf->data(), note_buf->size())) {
    PLOG_ERROR << "Failed to read PT_NOTE segment";
    return EX_IOERR;
  }

  return EX_OK;
}

bool CoredumpWriter::GetFileMappings(const std::vector<char> &note_buf,
                                     FileMappings *file_mappings) {
  // Locate NT_FILE note.
  ElfCoreDump::Note note({ note_buf.data(), note_buf.size() });
  while (note.IsValid() && note.GetType() != NT_FILE)
    note = note.GetNextNote();

  if (!note.IsValid()) {
    LOG_ERROR << "Failed to locate NT_FILE note";
    return false;
  }
  const auto note_desc = note.GetDescription();

  // The NT_FILE note is generated by the kernel's fs/binfmt_elf.c as a
  // sequence of long values, followed by a sequence of char values.
  struct FileNote {
    ElfW(Off) count;
    ElfW(Off) page_size;

    struct {
      ElfW(Addr) start;
      ElfW(Addr) end;
      ElfW(Off) offset;
    } files[1];

    // The |files| array has |count| elements. The file paths are stored after
    // the array, as |count| null-terminated strings.
  };
  const FileNote *file_note =
      reinterpret_cast<const FileNote *>(note_desc.data());
  const size_t num_files = file_note->count;
  const char *path =
      reinterpret_cast<const char *>(&file_note->files[num_files]);

  // Populate file mappings.
  for (size_t i = 0; i < num_files; ++i) {
    const auto file = &file_note->files[i];
    const auto it = file_mappings->insert({
        FileRange(file->start, file->end),
        { file->offset * file_note->page_size, path }
    });
    // Skip past NUL to next path.
    path += it.first->second.path.length() + 1;
  }
  // The last path should end the note.
  if (path != reinterpret_cast<const char *>(file_note) + note_desc.length()) {
    LOG_ERROR << "Invalid NT_FILE note";
    return false;
  }
  return true;
}

void CoredumpWriter::StripSegments(
    const std::vector<Phdr> &program_headers,
    const FileMappings &file_mappings,
    std::vector<Phdr> *stripped_program_headers) {
  stripped_program_headers->resize(program_headers.size());

  // The first segment has type PT_NOTE. Use the original data unchanged.
  (*stripped_program_headers)[0] = program_headers[0];

  for (size_t i = 1; i < program_headers.size(); ++i) {
    Phdr &out = (*stripped_program_headers)[i];
    out = program_headers[i];

    // If the type is PT_LOAD and the range is found in the set, the segment is
    // backed by a file, so it can be excluded as it doesn't contain stack data
    // useful to generate a minidump.
    const FileRange range(out.p_vaddr, out.p_vaddr + out.p_memsz);
    if (out.p_type == PT_LOAD && file_mappings.count(range))
      out.p_filesz = 0;

    // Calculate offset.
    const Phdr &prev_program_header = (*stripped_program_headers)[i - 1];
    out.p_offset = prev_program_header.p_offset + prev_program_header.p_filesz;
    // Offset alignment.
    if (out.p_align != 0 && out.p_offset % out.p_align != 0)
      out.p_offset += out.p_align - out.p_offset % out.p_align;
  }
}

int CoredumpWriter::WriteAuxv(const std::vector<char> &note_buf) {
  // Locate NT_AUXV note.
  ElfCoreDump::Note note({ note_buf.data(), note_buf.size() });
  while (note.IsValid() && note.GetType() != NT_AUXV)
    note = note.GetNextNote();

  if (!note.IsValid()) {
    LOG_ERROR << "Failed to locate NT_AUXV note";
    return EX_OSFILE;
  }

  const auto path = container_dir_ + std::string("/auxv");
  const ScopedFd auxv(TEMP_FAILURE_RETRY(
      open(path.c_str(), O_WRONLY | O_CREAT | O_EXCL, S_IRUSR | S_IWUSR)));

  if (!auxv.IsValid()) {
    PLOG_ERROR << "Failed to open '" << path.c_str() << "'";
    return EX_CANTCREAT;
  }

  // The NT_AUXV note has the same format as /proc/[pid]/auxv.
  const auto desc = note.GetDescription();
  return WriteAllBlocking(auxv, desc.data(), desc.length()) ? EX_OK : EX_IOERR;
}

int CoredumpWriter::WriteMaps(const std::vector<Phdr> &program_headers,
                              const FileMappings &file_mappings) {
  const auto path = container_dir_ + std::string("/maps");
  const ScopedFd maps(TEMP_FAILURE_RETRY(
      open(path.c_str(), O_WRONLY | O_CREAT | O_EXCL, S_IRUSR | S_IWUSR)));

  if (!maps.IsValid()) {
    PLOG_ERROR << "Failed to open '" << path.c_str() << "'";
    return EX_CANTCREAT;
  }

  for (const auto &program_header : program_headers) {
    if (program_header.p_type != PT_LOAD)
      continue;
    const FileRange range(program_header.p_vaddr,
                          program_header.p_vaddr + program_header.p_memsz);
    // If a mapping is found for the range, the range is mapped to a file.
    const auto it = file_mappings.find(range);
    const ElfW(Off) offset = it != file_mappings.end() ? it->second.offset : 0;
    const auto path = it != file_mappings.end() ? it->second.path.c_str() : "";

    const int kBufSize = 1024;
    char buf[kBufSize];
    // See kernel's fs/proc/task_{no,}mmu.c for format string.
    const int len = snprintf(
        buf, kBufSize,
        "%08" PRIxPTR
        "-"
        "%08" PRIxPTR
        " %c%c%c%c"
        " %08" PRIx64
        " %02x:%02x %lu %s\n",
        range.first, range.second,
        program_header.p_flags & PF_R ? 'r' : '-',
        program_header.p_flags & PF_W ? 'w' : '-',
        program_header.p_flags & PF_X ? 'x' : '-',
        'p',  // Fake value: we can't know if the mapping is shared or private.
        static_cast<uint64_t>(offset),
        0,  // Fake device (major) value.
        0,  // Fake device (minor) value.
        0ul,  // Fake inode value.
        path);
    if (len < 0 || len >= kBufSize || !WriteAllBlocking(maps, buf, len))
      return EX_OSFILE;
  }

  return EX_OK;
}