// Copyright 2019 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 <inttypes.h>

#include <ddk/protocol/block.h>
#include <ddk/protocol/block/volume.h>
#include <fbl/unique_ptr.h>
#include <lib/sync/completion.h>
#include <lib/zircon-internal/debug.h>
#include <lib/zx/vmo.h>
#include <zircon/status.h>
#include <zxcrypt/ddk-volume.h>
#include <zxcrypt/volume.h>

#include <memory>
#include <utility>

#define ZXDEBUG 0

namespace zxcrypt {

void SyncComplete(void* cookie, zx_status_t status, block_op_t* block) {
    // Use the 32bit command field to shuttle the response back to the callsite that's waiting on
    // the completion
    block->command = status;
    sync_completion_signal(static_cast<sync_completion_t*>(cookie));
}

// Performs synchronous I/O
zx_status_t SyncIO(zx_device_t* dev, uint32_t cmd, void* buf, size_t off, size_t len) {
    zx_status_t rc;

    if (!dev || !buf || len == 0) {
        xprintf("bad parameter(s): dev=%p, buf=%p, len=%zu\n", dev, buf, len);
        return ZX_ERR_INVALID_ARGS;
    }

    block_impl_protocol_t proto;
    if ((rc = device_get_protocol(dev, ZX_PROTOCOL_BLOCK, &proto)) != ZX_OK) {
        xprintf("block protocol not support\n");
        return ZX_ERR_NOT_SUPPORTED;
    }

    zx::vmo vmo;
    if ((rc = zx::vmo::create(len, 0, &vmo)) != ZX_OK) {
        xprintf("zx::vmo::create failed: %s\n", zx_status_get_string(rc));
        return rc;
    }

    block_info_t info;
    size_t op_size;
    block_impl_query(&proto, &info, &op_size);

    size_t bsz = info.block_size;
    ZX_DEBUG_ASSERT(off / bsz <= UINT32_MAX);
    ZX_DEBUG_ASSERT(len / bsz <= UINT32_MAX);
    fbl::AllocChecker ac;
    std::unique_ptr<char[]> raw;
    if constexpr (alignof(block_op_t) > __STDCPP_DEFAULT_NEW_ALIGNMENT__) {
        raw = std::unique_ptr<char[]>(
            new (static_cast<std::align_val_t>(alignof(block_op_t)), &ac) char[op_size]);
    } else {
        raw = std::unique_ptr<char[]>(
            new (&ac) char[op_size]);
    }
    if (!ac.check()) {
        return ZX_ERR_NO_MEMORY;
    }
    block_op_t* block = reinterpret_cast<block_op_t*>(raw.get());

    sync_completion_t completion;
    sync_completion_reset(&completion);

    block->command = cmd;
    block->rw.vmo = vmo.get();
    block->rw.length = static_cast<uint32_t>(len / bsz);
    block->rw.offset_dev = static_cast<uint32_t>(off / bsz);
    block->rw.offset_vmo = 0;

    if (cmd == BLOCK_OP_WRITE && (rc = vmo.write(buf, 0, len)) != ZX_OK) {
        xprintf("zx::vmo::write failed: %s\n", zx_status_get_string(rc));
        return rc;
    }

    block_impl_queue(&proto, block, SyncComplete, &completion);
    sync_completion_wait(&completion, ZX_TIME_INFINITE);

    rc = block->command;
    if (rc != ZX_OK) {
        xprintf("Block I/O failed: %s\n", zx_status_get_string(rc));
        return rc;
    }

    if (cmd == BLOCK_OP_READ && (rc = vmo.read(buf, 0, len)) != ZX_OK) {
        xprintf("zx::vmo::read failed: %s\n", zx_status_get_string(rc));
        return rc;
    }

    return ZX_OK;
}

DdkVolume::DdkVolume(zx_device_t* dev) : Volume(), dev_(dev) {}

zx_status_t DdkVolume::Bind(crypto::Cipher::Direction direction, crypto::Cipher* cipher) const {
    zx_status_t rc;
    ZX_DEBUG_ASSERT(dev_); // Cannot bind from library

    if (!cipher) {
        xprintf("bad parameter(s): cipher=%p\n", cipher);
        return ZX_ERR_INVALID_ARGS;
    }
    if (!block_.get()) {
        xprintf("not initialized\n");
        return ZX_ERR_BAD_STATE;
    }
    if ((rc = cipher->Init(cipher_, direction, data_key_, data_iv_, block_.len())) != ZX_OK) {
        return rc;
    }

    return ZX_OK;
}

zx_status_t DdkVolume::Unlock(zx_device_t* dev, const crypto::Secret& key, key_slot_t slot,
                                       fbl::unique_ptr<DdkVolume>* out) {
    zx_status_t rc;

    if (!dev || !out) {
        xprintf("bad parameter(s): dev=%p, out=%p\n", dev, out);
        return ZX_ERR_INVALID_ARGS;
    }
    fbl::AllocChecker ac;
    fbl::unique_ptr<DdkVolume> volume(new (&ac) DdkVolume(dev));
    if (!ac.check()) {
        xprintf("allocation failed: %zu bytes\n", sizeof(DdkVolume));
        return ZX_ERR_NO_MEMORY;
    }
    if ((rc = volume->Init()) != ZX_OK) {
        xprintf("volume->Init() failed: %s\n", zx_status_get_string(rc));
        return rc;
    }
    if ((rc = volume->Unlock(key, slot)) != ZX_OK) {
        xprintf("volume->Unlock() failed: %s\n", zx_status_get_string(rc));
        return rc;
    }

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

zx_status_t DdkVolume::Init() {
    return Volume::Init();
}

zx_status_t DdkVolume::Unlock(const crypto::Secret& key, key_slot_t slot) {
    return Volume::Unlock(key, slot);
}

zx_status_t DdkVolume::GetBlockInfo(BlockInfo* out) {
    zx_status_t rc;
    block_impl_protocol_t proto;
    if ((rc = device_get_protocol(dev_, ZX_PROTOCOL_BLOCK, &proto)) != ZX_OK) {
        xprintf("block protocol not supported\n");
        return ZX_ERR_NOT_SUPPORTED;
    }

    block_info_t info;
    size_t block_op_size;
    block_impl_query(&proto, &info, &block_op_size);

    out->block_count = info.block_count;
    out->block_size = info.block_size;
    return ZX_OK;
}

zx_status_t DdkVolume::GetFvmSliceSize(uint64_t* out) {
    zx_status_t rc;
    block_volume_protocol_t proto;
    if ((rc = device_get_protocol(dev_, ZX_PROTOCOL_BLOCK_VOLUME, &proto)) != ZX_OK) {
        xprintf("block volume protocol not supported\n");
        return ZX_ERR_NOT_SUPPORTED;
    }

    parent_volume_info_t info;
    if ((rc = block_volume_query(&proto, &info)) != ZX_OK) {
        return rc;
    }

    *out = info.slice_size;
    return ZX_OK;
}

zx_status_t DdkVolume::DoBlockFvmVsliceQuery(uint64_t vslice_start,
                                             SliceRegion ranges[MAX_SLICE_REGIONS],
                                             uint64_t* slice_count) {
    static_assert(MAX_SLICE_QUERY_REQUESTS == Volume::MAX_SLICE_REGIONS,
                  "block volume slice response count must match");
    zx_status_t rc;
    block_volume_protocol_t proto;
    if ((rc = device_get_protocol(dev_, ZX_PROTOCOL_BLOCK_VOLUME, &proto)) != ZX_OK) {
        xprintf("block volume protocol not supported\n");
        return ZX_ERR_NOT_SUPPORTED;
    }

    slice_region_t tmp_ranges[MAX_SLICE_QUERY_REQUESTS];
    size_t range_count;
    if ((rc = block_volume_query_slices(&proto, &vslice_start, 1, tmp_ranges,
                                        MAX_SLICE_QUERY_REQUESTS, &range_count)) != ZX_OK) {
        return rc;
    }
    if (range_count > Volume::MAX_SLICE_REGIONS) {
        // Should be impossible.  Trust nothing.
        return ZX_ERR_BAD_STATE;
    }

    *slice_count = range_count;
    for (size_t i = 0; i < range_count; i++) {
        ranges[i].allocated = tmp_ranges[i].allocated;
        ranges[i].count = tmp_ranges[i].count;
    }
    return ZX_OK;
}

zx_status_t DdkVolume::DoBlockFvmExtend(uint64_t start_slice, uint64_t slice_count) {
    zx_status_t rc;
    block_volume_protocol_t proto;
    if ((rc = device_get_protocol(dev_, ZX_PROTOCOL_BLOCK_VOLUME, &proto)) != ZX_OK) {
        xprintf("block volume protocol not supported\n");
        return ZX_ERR_NOT_SUPPORTED;
    }

    slice_extent_t extent;
    extent.offset = start_slice;
    extent.length = slice_count;
    return block_volume_extend(&proto, &extent);
}

zx_status_t DdkVolume::Read() {
    return SyncIO(dev_, BLOCK_OP_READ, block_.get(), offset_, block_.len());
}

zx_status_t DdkVolume::Write() {
    return SyncIO(dev_, BLOCK_OP_WRITE, block_.get(), offset_, block_.len());
}

} // namespace zxcrypt