// Copyright 2017 The Fuchsia 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 "fvm/sparse-reader.h"

#include <cinttypes>
#include <utility>

namespace fvm {

using Buffer = internal::Buffer;

Buffer::Buffer() = default;
Buffer::Buffer(uint64_t offset, size_t size, uint64_t capacity)
    : capacity_(capacity) {
    info_.size = size;
    info_.offset = offset;
    data_.reset(new uint8_t[capacity]);
}
Buffer::Buffer(Buffer&&) = default;
Buffer& Buffer::operator=(Buffer&&) = default;
Buffer::~Buffer() = default;

bool Buffer::IsEmpty() const {
    return info_.offset == 0 && info_.size == 0;
}

void Buffer::Write(uint8_t* data, size_t length) {
    ZX_ASSERT(length <= capacity_);
    // We should have read all previous data from buffer before writing more.
    ZX_ASSERT(IsEmpty());

    if (length > 0) {
        memcpy(data_.get(), data, length);
        info_.size = length;
    }
}

void Buffer::Read(uint8_t* target, size_t length, size_t* actual) {
    size_t cp_sz = fbl::min(length, info_.size);

    if (cp_sz > 0) {
        memcpy(target, data_.get() + info_.offset, cp_sz);
        info_.offset += cp_sz;
    }

    info_.size -= cp_sz;

    if (info_.size == 0) {
        info_.offset = 0;
    }

    *actual = cp_sz;
}

zx_status_t SparseReader::Create(fbl::unique_fd fd, fbl::unique_ptr<SparseReader>* out) {
    return SparseReader::CreateHelper(std::move(fd), true /* verbose */, out);
}
zx_status_t SparseReader::CreateSilent(fbl::unique_fd fd, fbl::unique_ptr<SparseReader>* out) {
    return SparseReader::CreateHelper(std::move(fd), false /* verbose */, out);
}

zx_status_t SparseReader::CreateHelper(fbl::unique_fd fd, bool verbose,
                                       fbl::unique_ptr<SparseReader>* out) {
    fbl::unique_ptr<SparseReader> reader(new SparseReader(std::move(fd), verbose));

    zx_status_t status;
    if ((status = reader->ReadMetadata()) != ZX_OK) {
        return status;
    }

    *out = std::move(reader);
    return ZX_OK;
}

SparseReader::SparseReader(fbl::unique_fd fd, bool verbose)
    : compressed_(false), verbose_(verbose), fd_(std::move(fd)) {}

zx_status_t SparseReader::ReadMetadata() {
    // Read sparse image header.
    fvm::sparse_image_t image;
    if (read(fd_.get(), &image, sizeof(fvm::sparse_image_t)) != sizeof(fvm::sparse_image_t)) {
        fprintf(stderr, "failed to read the sparse header\n");
        return ZX_ERR_IO;
    }

    // Verify the header.
    if (image.magic != fvm::kSparseFormatMagic) {
        fprintf(stderr, "SparseReader: Bad magic\n");
        return ZX_ERR_BAD_STATE;
    } else if (image.version != fvm::kSparseFormatVersion) {
        fprintf(stderr, "SparseReader: Unexpected sparse file version\n");
        return ZX_ERR_BAD_STATE;
    }

    metadata_.reset(new uint8_t[image.header_length]);
    memcpy(metadata_.get(), &image, sizeof(image));

    // Read remainder of metadata.
    size_t off = sizeof(image);
    while (off < image.header_length) {
        ssize_t r = read(fd_.get(), &metadata_[off], image.header_length - off);
        if (r < 0) {
            fprintf(stderr, "SparseReader: Failed to read metadata\n");
            return ZX_ERR_IO;
        }
        off += r;
    }

    // If image is compressed, additional setup is required
    if (image.flags & fvm::kSparseFlagLz4) {
        if (verbose_) {
            printf("Found compressed file\n");
        }

        compressed_ = true;
        // Initialize decompression context
        LZ4F_errorCode_t errc = LZ4F_createDecompressionContext(&dctx_, LZ4F_VERSION);
        if (LZ4F_isError(errc)) {
            fprintf(stderr, "SparseReader: could not initialize decompression: %s\n",
                    LZ4F_getErrorName(errc));
            return ZX_ERR_INTERNAL;
        }

        size_t src_sz = 4;
        size_t dst_sz = 0;
        fbl::unique_ptr<uint8_t[]> inbufptr(new uint8_t[src_sz]);

        uint8_t* inbuf = inbufptr.get();

        // Read first 4 bytes to let LZ4 tell us how much it expects in the first pass.
        ssize_t nr = read(fd_.get(), inbuf, src_sz);
        if (nr < static_cast<ssize_t>(src_sz)) {
            fprintf(stderr, "SparseReader: could not read from input\n");
            return ZX_ERR_IO;
        }

        // Run decompress once to find out how much data we should read for the next decompress run
        // Since we are not yet decompressing any actual data, the dst_buffer is null
        to_read_ = LZ4F_decompress(dctx_, nullptr, &dst_sz, inbuf, &src_sz, NULL);
        if (LZ4F_isError(to_read_)) {
            fprintf(stderr, "SparseReader: could not decompress header: %s\n",
                    LZ4F_getErrorName(to_read_));
            return ZX_ERR_INTERNAL;
        }

        if (to_read_ > LZ4_MAX_BLOCK_SIZE) {
            to_read_ = LZ4_MAX_BLOCK_SIZE;
        }

        // Initialize data buffers
        zx_status_t status;
        if ((status = InitializeBuffer(LZ4_MAX_BLOCK_SIZE, &out_)) != ZX_OK) {
            return status;
        } else if ((status = InitializeBuffer(LZ4_MAX_BLOCK_SIZE, &in_)) != ZX_OK) {
            return status;
        }
    }

    return ZX_OK;
}

zx_status_t SparseReader::InitializeBuffer(size_t size, Buffer* out_buffer) {
    if (size < LZ4_MAX_BLOCK_SIZE) {
        fprintf(stderr, "Buffer size must be >= %d\n", LZ4_MAX_BLOCK_SIZE);
        return ZX_ERR_INVALID_ARGS;
    }
    *out_buffer = Buffer(0, 0, size);
    return ZX_OK;
}

SparseReader::~SparseReader() {
    PrintStats();

    if (compressed_) {
        LZ4F_freeDecompressionContext(dctx_);
    }
}

fvm::sparse_image_t* SparseReader::Image() {
    return reinterpret_cast<fvm::sparse_image_t*>(metadata_.get());
}

fvm::partition_descriptor_t* SparseReader::Partitions() {
    return reinterpret_cast<fvm::partition_descriptor_t*>(
        reinterpret_cast<uintptr_t>(metadata_.get()) + sizeof(fvm::sparse_image_t));
}

zx_status_t SparseReader::ReadData(uint8_t* data, size_t length, size_t* actual) {
#ifdef __Fuchsia__
    zx_ticks_t start = zx_ticks_get();
#endif
    size_t total_size = 0;
    if (compressed_) {
        if (out_.IsEmpty() && to_read_ == 0) {
            // There is no more to read
            return ZX_ERR_OUT_OF_RANGE;
        }

        // Read previously decompressed data from buffer if possible
        out_.Read(data, length, &total_size);

        // If we still have data to read, start decompression (reading more from fd as needed)
        while (total_size < length && to_read_ > 0) {
            // Make sure data to read does not exceed max, and both buffers are empty
            ZX_ASSERT(out_.IsEmpty());
            ZX_ASSERT(in_.IsEmpty());
            ZX_ASSERT(to_read_ <= in_.capacity());

            // Read specified amount from fd
            zx_status_t status;
            if ((status = ReadRaw(in_.get(), to_read_, &(in_.info()->size))) != ZX_OK) {
                return status;
            }

            size_t src_sz = in_.size();
            size_t next = 0;

            // Decompress all compressed data
            while (in_.offset() < to_read_) {
                size_t dst_sz = out_.capacity() - out_.size();
                next = LZ4F_decompress(dctx_, out_.get() + out_.size(), &dst_sz,
                                       in_.get() + in_.offset(), &src_sz, NULL);
                if (LZ4F_isError(next)) {
                    fprintf(stderr, "could not decompress input: %s\n", LZ4F_getErrorName(next));
                    return -1;
                }

                out_.info()->size += dst_sz;
                in_.info()->offset += src_sz;
                in_.info()->size -= src_sz;
                src_sz = to_read_ - in_.offset();
            }

            // Make sure we have read all data from in_buf_
            if (in_.size() > 0) {
                return ZX_ERR_IO;
            }

            in_.info()->offset = 0;

            // Copy newly decompressed data from outbuf
            size_t cp = fbl::min(length - total_size, static_cast<size_t>(out_.size()));
            out_.Read(data + total_size, cp, &cp);
            total_size += cp;
            to_read_ = next;

            if (to_read_ > LZ4_MAX_BLOCK_SIZE) {
                to_read_ = LZ4_MAX_BLOCK_SIZE;
            }
        }
    } else {
        zx_status_t status = ReadRaw(data, length, &total_size);

        if (status != ZX_OK) {
            return status;
        }
    }

#ifdef __Fuchsia__
    total_time_ += zx_ticks_get() - start;
#endif
    *actual = total_size;
    return ZX_OK;
}

zx_status_t SparseReader::ReadRaw(uint8_t* data, size_t length, size_t* actual) {
#ifdef __Fuchsia__
    zx_ticks_t start = zx_ticks_get();
#endif
    ssize_t r;
    size_t total_size = 0;
    size_t bytes_left = length;
    while ((r = read(fd_.get(), data + total_size, bytes_left)) > 0) {
        total_size += r;
        bytes_left -= r;
        if (bytes_left == 0) {
            break;
        }
    }

#ifdef __Fuchsia__
    read_time_ += zx_ticks_get() - start;
#endif

    if (r < 0) {
        return static_cast<zx_status_t>(r);
    }

    *actual = total_size;
    return ZX_OK;
}

zx_status_t SparseReader::WriteDecompressed(fbl::unique_fd outfd) {
    if (!compressed_) {
        fprintf(stderr, "BlockReader: File is not compressed\n");
        return ZX_ERR_INVALID_ARGS;
    }

    // Update metadata and write to new file.
    fvm::sparse_image_t* image = Image();
    image->flags &= ~fvm::kSparseFlagLz4;

    if (write(outfd.get(), metadata_.get(), image->header_length) !=
        static_cast<ssize_t>(image->header_length)) {
        fprintf(stderr, "BlockReader: could not write header to out file\n");
        return -1;
    }

    // Read/write decompressed data in LZ4_MAX_BLOCK_SIZE chunks.
    while (true) {
        zx_status_t status;
        uint8_t data[LZ4_MAX_BLOCK_SIZE];
        size_t length;
        if ((status = ReadData(data, LZ4_MAX_BLOCK_SIZE, &length)) != ZX_OK) {
            if (status == ZX_ERR_OUT_OF_RANGE) {
                return ZX_OK;
            }

            return status;
        }

        if (write(outfd.get(), data, length) != static_cast<ssize_t>(length)) {
            fprintf(stderr, "BlockReader: failed to write to output\n");
            return ZX_ERR_IO;
        }
    }
}

void SparseReader::PrintStats() const {
    if (verbose_) {
        printf("Reading FVM from compressed file: %s\n", compressed_ ? "true" : "false");
        printf("Remaining bytes read into compression buffer:    \%" PRIuMAX "\n", in_.size());
        printf("Remaining bytes written to decompression buffer:\%" PRIuMAX "\n", out_.size());
#ifdef __Fuchsia__
        printf("Time reading bytes from sparse FVM file:   %lu (%lu s)\n", read_time_,
               read_time_ / zx_ticks_per_second());
        printf("Time reading bytes AND decompressing them: %lu (%lu s)\n", total_time_,
               total_time_ / zx_ticks_per_second());
#endif
    }
}

} // namespace fvm