/*
 * Copyright (C) 2016-2024 Apple Inc. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY APPLE INC. ``AS IS'' AND ANY
 * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
 * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL APPLE INC. OR
 * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
 * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include "config.h"
#include "WasmMemory.h"

#if ENABLE(WEBASSEMBLY)

#include "JSCInlines.h"
#include "JSWebAssemblyInstance.h"
#include "Options.h"
#include "WasmFaultSignalHandler.h"
#include "WasmLimits.h"
#include "WeakGCSetInlines.h"
#include <bmalloc/Gigacage.h>
#include <wtf/CheckedArithmetic.h>
#include <wtf/DataLog.h>
#include <wtf/FastMalloc.h>
#include <wtf/Lock.h>
#include <wtf/Platform.h>
#include <wtf/PrintStream.h>
#include <wtf/TZoneMallocInlines.h>

#include <cstring>

WTF_ALLOW_UNSAFE_BUFFER_USAGE_BEGIN

namespace JSC {

namespace Wasm {

WTF_MAKE_TZONE_ALLOCATED_IMPL(Memory);

// FIXME: We could be smarter about memset / mmap / madvise. https://bugs.webkit.org/show_bug.cgi?id=170343
// FIXME: Give up some of the cached fast memories if the GC determines it's easy to get them back, and they haven't been used in a while. https://bugs.webkit.org/show_bug.cgi?id=170773
// FIXME: Limit slow memory size. https://bugs.webkit.org/show_bug.cgi?id=170825

namespace {

constexpr bool verbose = false;

NEVER_INLINE NO_RETURN_DUE_TO_CRASH void webAssemblyCouldntGetFastMemory() { CRASH(); }

// A grow's resulting page count must be one the memory's own address type can declare, which is
// tighter than what PageCount can represent.
static bool isDeclarablePageCount(PageCount pageCount, AddressType addressType)
{
    return pageCount && pageCount.isValid() && pageCount.pageCount() <= maxDeclarablePages(addressType);
}

template<typename Func>
static bool tryAllocate(VM& vm, const Func& allocate)
{
    unsigned numTries = 2;
    bool done = false;
    for (unsigned i = 0; i < numTries && !done; ++i) {
        switch (allocate()) {
        case BufferMemoryResult::Kind::Success:
            done = true;
            break;
        case BufferMemoryResult::Kind::SuccessAndNotifyMemoryPressure:
            vm.heap.collectAsync(CollectionScope::Full);
            done = true;
            break;
        case BufferMemoryResult::Kind::SyncTryToReclaimMemory:
            if (i + 1 == numTries)
                break;
            vm.heap.collectSync(CollectionScope::Full);
            break;
        }
    }
    return done;
}

} // anonymous namespace

uint64_t maxAllocatableBytes(AddressType addressType)
{
    return std::min<uint64_t>(maxBufferByteLength(addressType), maxGrowableBufferReservationBytes);
}

Memory::Memory()
    : m_handle(adoptRef(*new BufferMemoryHandle(BufferMemoryHandle::nullBasePointer(), 0, 0, PageCount(0), PageCount(0), MemorySharingMode::Default, MemoryMode::BoundsChecking)))
{
}

Memory::Memory(PageCount initial, PageCount maximum, MemorySharingMode sharingMode, AddressType addressType, WTF::Function<void(GrowSuccess, PageCount, PageCount)>&& growSuccessCallback)
    : m_handle(adoptRef(*new BufferMemoryHandle(BufferMemoryHandle::nullBasePointer(), 0, 0, initial, maximum, sharingMode, MemoryMode::BoundsChecking)))
    , m_growSuccessCallback(WTF::move(growSuccessCallback))
    , m_addressType(addressType)
{
    ASSERT(!initial.bytes());
    ASSERT(mode() == MemoryMode::BoundsChecking);
    dataLogLnIf(verbose, "Memory::Memory allocating ", *this);
    ASSERT(basePointer());
}

Memory::Memory(Ref<BufferMemoryHandle>&& handle, AddressType addressType, WTF::Function<void(GrowSuccess, PageCount, PageCount)>&& growSuccessCallback)
    : m_handle(WTF::move(handle))
    , m_growSuccessCallback(WTF::move(growSuccessCallback))
    , m_addressType(addressType)
{
    dataLogLnIf(verbose, "Memory::Memory allocating ", *this);
}

Memory::Memory(Ref<BufferMemoryHandle>&& handle, Ref<SharedArrayBufferContents>&& shared, AddressType addressType, WTF::Function<void(GrowSuccess, PageCount, PageCount)>&& growSuccessCallback)
    : m_handle(WTF::move(handle))
    , m_shared(WTF::move(shared))
    , m_growSuccessCallback(WTF::move(growSuccessCallback))
    , m_addressType(addressType)
{
    dataLogLnIf(verbose, "Memory::Memory allocating ", *this);
}

Ref<Memory> Memory::create()
{
    return adoptRef(*new Memory());
}

Ref<Memory> Memory::create(Ref<BufferMemoryHandle>&& handle, AddressType addressType, WTF::Function<void(GrowSuccess, PageCount, PageCount)>&& growSuccessCallback)
{
    return adoptRef(*new Memory(WTF::move(handle), addressType, WTF::move(growSuccessCallback)));
}

Ref<Memory> Memory::create(Ref<SharedArrayBufferContents>&& shared, AddressType addressType, WTF::Function<void(GrowSuccess, PageCount, PageCount)>&& growSuccessCallback)
{
    RefPtr<BufferMemoryHandle> handle = shared->memoryHandle();
    ASSERT(handle);
    return adoptRef(*new Memory(handle.releaseNonNull(), WTF::move(shared), addressType, WTF::move(growSuccessCallback)));
}

Ref<Memory> Memory::createZeroSized(MemorySharingMode sharingMode, AddressType addressType, WTF::Function<void(GrowSuccess, PageCount, PageCount)>&& growSuccessCallback)
{
    return adoptRef(*new Memory(PageCount(0), PageCount(0), sharingMode, addressType, WTF::move(growSuccessCallback)));
}

RefPtr<Memory> Memory::tryCreate(VM& vm, PageCount initial, PageCount maximum, MemorySharingMode sharingMode, AddressType addressType, std::optional<MemoryMode> desiredMemoryMode, WTF::Function<void(GrowSuccess, PageCount, PageCount)>&& growSuccessCallback)
{
    ASSERT(initial);
    RELEASE_ASSERT(!maximum || maximum >= initial); // This should be guaranteed by our caller.

    const uint64_t initialBytes = initial.bytes();
    const uint64_t allocatableBytes = maxAllocatableBytes(addressType);
    const uint64_t declaredMaximumBytes = maximum ? maximum.bytes() : 0;
    // A buffer must not advertise a length this platform could never grow the memory to, so what it
    // reports is bounded by what can be allocated rather than by what the address type may declare.
    const size_t reservedMaximumBytes = std::min(declaredMaximumBytes, allocatableBytes);

    if (initialBytes > allocatableBytes)
        return nullptr; // Client will throw OOMError.

    if (maximum && !declaredMaximumBytes) {
        // User specified a zero maximum, initial size must also be zero.
        RELEASE_ASSERT(!initialBytes);
        return createZeroSized(sharingMode, addressType, WTF::move(growSuccessCallback));
    }

    bool done = tryAllocate(vm,
        [&] () -> BufferMemoryResult::Kind {
            return BufferMemoryManager::singleton().tryAllocatePhysicalBytes(initialBytes);
        });
    if (!done)
        return nullptr;

    char* fastMemory = nullptr;
    if (Options::useWasmFastMemory() && desiredMemoryMode.value_or(MemoryMode::Signaling) == MemoryMode::Signaling && !addressType.is64Bit()) {
#if CPU(ADDRESS32)
        RELEASE_ASSERT_NOT_REACHED_WITH_MESSAGE("32-bit platforms don't support fast memory.");
#endif
        tryAllocate(vm,
            [&] () -> BufferMemoryResult::Kind {
                auto result = BufferMemoryManager::singleton().tryAllocateFastMemory();
                fastMemory = std::bit_cast<char*>(result.basePtr);
                return result.kind;
            });
    }
    
    if (fastMemory) {
        constexpr bool readable = false;
        constexpr bool writable = false;
        OSAllocator::protect(fastMemory + initialBytes, BufferMemoryHandle::fastMappedBytes() - initialBytes, readable, writable);
        switch (sharingMode) {
        case MemorySharingMode::Default: {
            return Memory::create(adoptRef(*new BufferMemoryHandle(fastMemory, initialBytes, BufferMemoryHandle::fastMappedBytes(), initial, maximum, MemorySharingMode::Default, MemoryMode::Signaling)), addressType, WTF::move(growSuccessCallback));
        }
        case MemorySharingMode::Shared: {
            Ref handle = adoptRef(*new BufferMemoryHandle(fastMemory, initialBytes, BufferMemoryHandle::fastMappedBytes(), initial, maximum, MemorySharingMode::Shared, MemoryMode::Signaling));
            auto span = handle->mutableSpan();
            auto content = SharedArrayBufferContents::create(span, reservedMaximumBytes, WTF::move(handle), nullptr, SharedArrayBufferContents::Mode::WebAssembly);
            return Memory::create(WTF::move(content), addressType, WTF::move(growSuccessCallback));
        }
        }
        RELEASE_ASSERT_NOT_REACHED();
        return nullptr;
    }

    if (desiredMemoryMode == MemoryMode::Signaling) {
        BufferMemoryManager::singleton().freePhysicalBytes(initialBytes);
        return nullptr;
    }

    if (Options::crashIfWasmCantFastMemory()) [[unlikely]]
        webAssemblyCouldntGetFastMemory();

    switch (sharingMode) {
    case MemorySharingMode::Default: {
        if (!initialBytes)
            return adoptRef(new Memory(initial, maximum, MemorySharingMode::Default, addressType, WTF::move(growSuccessCallback)));

        void* slowMemory = Gigacage::tryAllocateZeroedVirtualPages(Gigacage::Primitive, initialBytes);
        if (!slowMemory) {
            BufferMemoryManager::singleton().freePhysicalBytes(initialBytes);
            return nullptr;
        }
        return Memory::create(adoptRef(*new BufferMemoryHandle(slowMemory, initialBytes, initialBytes, initial, maximum, MemorySharingMode::Default, MemoryMode::BoundsChecking)), addressType, WTF::move(growSuccessCallback));
    }
    case MemorySharingMode::Shared: {
        char* slowMemory = nullptr;
        tryAllocate(vm,
            [&] () -> BufferMemoryResult::Kind {
                auto result = BufferMemoryManager::singleton().tryAllocateGrowableBoundsCheckingMemory(reservedMaximumBytes);
                slowMemory = std::bit_cast<char*>(result.basePtr);
                return result.kind;
            });
        if (!slowMemory) {
            BufferMemoryManager::singleton().freePhysicalBytes(initialBytes);
            return nullptr;
        }

        constexpr bool readable = false;
        constexpr bool writable = false;
        OSAllocator::protect(slowMemory + initialBytes, reservedMaximumBytes - initialBytes, readable, writable);

        Ref handle = adoptRef(*new BufferMemoryHandle(slowMemory, initialBytes, reservedMaximumBytes, initial, maximum, MemorySharingMode::Shared, MemoryMode::BoundsChecking));
        auto span = handle->mutableSpan();
        auto content = SharedArrayBufferContents::create(span, reservedMaximumBytes, WTF::move(handle), nullptr, SharedArrayBufferContents::Mode::WebAssembly);
        return Memory::create(WTF::move(content), addressType, WTF::move(growSuccessCallback));
    }
    }
    RELEASE_ASSERT_NOT_REACHED();
    return nullptr;
}

Memory::~Memory() = default;

bool Memory::addressIsInGrowableOrFastMemory(void* address)
{
    return BufferMemoryManager::singleton().isInGrowableOrFastMemory(address);
}

Expected<PageCount, GrowFailReason> Memory::growShared(VM& vm, PageCount delta)
{
    // Collections happen with the handle lock released, so a retry recomputes the target: another agent
    // may have grown this memory while we were unlocked, and delta is relative to whatever the size is now.
    constexpr unsigned maximumAttempts = 2;
    for (unsigned attempt = 1; ; ++attempt) {
        PageCount oldPageCount;
        PageCount newPageCount;
        Expected<int64_t, GrowFailReason> result;
        auto collection = BufferMemoryResult::Kind::Success;
        {
            std::optional<Locker<Lock>> locker;
            // m_shared may not be exist, if this is zero byte memory with zero byte maximum size.
            if (m_shared)
                locker.emplace(m_shared->memoryHandle()->lock());

            oldPageCount = PageCount::fromBytes(size());
            newPageCount = oldPageCount + delta;
            if (!isDeclarablePageCount(newPageCount, addressType()))
                return makeUnexpected(GrowFailReason::InvalidGrowSize);
            const uint64_t allocatableBytes = maxAllocatableBytes(addressType());
            if (newPageCount.bytes() > allocatableBytes)
                return makeUnexpected(GrowFailReason::OutOfMemory);

            if (!delta.pageCount())
                return oldPageCount;

            dataLogLnIf(verbose, "Memory::grow(", delta, ") to ", newPageCount, " from ", *this);
            RELEASE_ASSERT(newPageCount > PageCount::fromBytes(size()));

            if (maximum() && newPageCount > maximum())
                return makeUnexpected(GrowFailReason::WouldExceedMaximum);

            size_t desiredSize = newPageCount.bytes();
            RELEASE_ASSERT(m_shared);
            RELEASE_ASSERT(desiredSize <= allocatableBytes);
            RELEASE_ASSERT(desiredSize > size());
            constexpr bool sizeMustBePageMultiple = true;
            result = m_shared->tryGrow(locker.value(), desiredSize, sizeMustBePageMultiple, collection);
        }
        if (collection == BufferMemoryResult::Kind::SyncTryToReclaimMemory) {
            if (attempt == maximumAttempts)
                return makeUnexpected(GrowFailReason::OutOfMemory);
            vm.heap.collectSync(CollectionScope::Full);
            continue;
        }
        if (!result)
            return makeUnexpected(result.error());

        // The instances' cached memories are brought up to date before anything can collect.
        m_growSuccessCallback(GrowSuccessTag, oldPageCount, newPageCount);
        if (collection == BufferMemoryResult::Kind::SuccessAndNotifyMemoryPressure)
            vm.heap.collectAsync(CollectionScope::Full);
        return oldPageCount;
    }
}

Expected<PageCount, GrowFailReason> Memory::grow(VM& vm, PageCount delta)
{
    if (!delta.isValid())
        return makeUnexpected(GrowFailReason::InvalidDelta);

    if (sharingMode() == MemorySharingMode::Shared)
        return growShared(vm, delta);

    ASSERT(!m_shared);
    const PageCount oldPageCount = PageCount::fromBytes(size());
    const PageCount newPageCount = oldPageCount + delta;
    if (!isDeclarablePageCount(newPageCount, addressType()))
        return makeUnexpected(GrowFailReason::InvalidGrowSize);
    const uint64_t allocatableBytes = maxAllocatableBytes(addressType());
    if (newPageCount.bytes() > allocatableBytes)
        return makeUnexpected(GrowFailReason::OutOfMemory);

    auto success = [&] () {
        // Update cache for instance
        {
            Locker locker { m_handle->lock() };
            for (Ref anchor : m_handle->anchors(locker)) {
                Locker locker { anchor->m_lock };
                if (JSWebAssemblyInstance* instance = anchor->instance())
                    instance->updateCachedMemories();
            }
        }
        m_growSuccessCallback(GrowSuccessTag, oldPageCount, newPageCount);
        return oldPageCount;
    };

    if (delta.pageCount() == 0)
        return success();

    dataLogLnIf(verbose, "Memory::grow(", delta, ") to ", newPageCount, " from ", *this);
    RELEASE_ASSERT(newPageCount > PageCount::fromBytes(size()));

    if (maximum() && newPageCount > maximum())
        return makeUnexpected(GrowFailReason::WouldExceedMaximum);

    size_t desiredSize = newPageCount.bytes();
    RELEASE_ASSERT(desiredSize <= allocatableBytes);
    RELEASE_ASSERT(desiredSize > size());
    switch (mode()) {
    case MemoryMode::BoundsChecking: {
        bool allocationSuccess = tryAllocate(vm,
            [&] () -> BufferMemoryResult::Kind {
                return BufferMemoryManager::singleton().tryAllocatePhysicalBytes(desiredSize);
            });
        if (!allocationSuccess)
            return makeUnexpected(GrowFailReason::OutOfMemory);

        RELEASE_ASSERT(maximum().bytes() != 0);

        void* newMemory = Gigacage::tryAllocateZeroedVirtualPages(Gigacage::Primitive, desiredSize);
        if (!newMemory) {
            BufferMemoryManager::singleton().freePhysicalBytes(desiredSize);
            return makeUnexpected(GrowFailReason::OutOfMemory);
        }

        memcpySpan(unsafeMakeSpan(static_cast<uint8_t*>(newMemory), desiredSize), m_handle->mutableSpan());
        Ref newHandle = adoptRef(*new BufferMemoryHandle(newMemory, desiredSize, desiredSize, initial(), maximum(), sharingMode(), MemoryMode::BoundsChecking));
        m_handle->transferAnchors(newHandle.get());
        m_handle = WTF::move(newHandle);

        ASSERT(basePointer() == newMemory);
        return success();
    }
    case MemoryMode::Signaling: {
        size_t extraBytes = desiredSize - size();
        RELEASE_ASSERT(extraBytes);
        bool allocationSuccess = tryAllocate(vm,
            [&] () -> BufferMemoryResult::Kind {
                return BufferMemoryManager::singleton().tryAllocatePhysicalBytes(extraBytes);
            });
        if (!allocationSuccess)
            return makeUnexpected(GrowFailReason::OutOfMemory);

        void* memory = this->basePointer();
        RELEASE_ASSERT(memory);

        // Signaling memory must have been pre-allocated virtually.
        uint8_t* startAddress = static_cast<uint8_t*>(memory) + size();
        
        dataLogLnIf(verbose, "Marking WebAssembly memory's ", RawPointer(memory), " as read+write in range [", RawPointer(startAddress), ", ", RawPointer(startAddress + extraBytes), ")");
        constexpr bool readable = true;
        constexpr bool writable = true;
        OSAllocator::protect(startAddress, extraBytes, readable, writable);
        m_handle->updateSize(desiredSize);
        return success();
    }
    }

    RELEASE_ASSERT_NOT_REACHED();
    return oldPageCount;
}

bool Memory::init(uint64_t offset, const uint8_t* data, uint32_t length)
{
    if (sumOverflows<uint64_t>(offset, length))
        return false;

    if (offset + length > m_handle->size())
        return false;

    if (!length)
        return true;

    memcpy(static_cast<uint8_t*>(basePointer()) + offset, data, length);
    return true;
}

void Memory::registerInstance(JSWebAssemblyInstance& instance)
{
    m_handle->registerInstance(instance);
}

void Memory::dump(PrintStream& out) const
{
    auto handle = m_handle.copyRef();
    out.print("Memory at ", RawPointer(handle->memory()), ", size ", handle->size(), "B capacity ", handle->mappedCapacity(), "B, initial ", handle->initial(), " maximum ", handle->maximum(), " mode ", handle->mode(), " sharingMode ", handle->sharingMode());
}

} // namespace Wasm

} // namespace JSC

WTF_ALLOW_UNSAFE_BUFFER_USAGE_END

#endif // ENABLE(WEBASSEMBLY)