Atmosphere/libraries/libmesosphere/include/mesosphere/kern_k_handle_table.hpp

/*
 * Copyright (c) 2018-2020 Atmosphère-NX
 *
 * This program is free software; you can redistribute it and/or modify it
 * under the terms and conditions of the GNU General Public License,
 * version 2, as published by the Free Software Foundation.
 *
 * This program is distributed in the hope it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
 * more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
#pragma once
#include <mesosphere/kern_common.hpp>
#include <mesosphere/kern_k_auto_object.hpp>
#include <mesosphere/kern_k_spin_lock.hpp>
#include <mesosphere/kern_k_thread.hpp>
#include <mesosphere/kern_k_interrupt_event.hpp>

namespace ams::kern {

    constexpr ALWAYS_INLINE util::BitPack32 GetHandleBitPack(ams::svc::Handle handle) {
        return util::BitPack32{handle};
    }

    class KProcess;
    class KThread;

    class KHandleTable {
        NON_COPYABLE(KHandleTable);
        NON_MOVEABLE(KHandleTable);
        public:
            static constexpr size_t MaxTableSize = 1024;
        private:
            using HandleRawValue = util::BitPack32::Field<0, BITSIZEOF(u32), u32>;
            using HandleEncoded  = util::BitPack32::Field<0, BITSIZEOF(ams::svc::Handle), ams::svc::Handle>;

            using HandleIndex    = util::BitPack32::Field<0,                    15, u16>;
            using HandleLinearId = util::BitPack32::Field<HandleIndex::Next,    15, u16>;
            using HandleReserved = util::BitPack32::Field<HandleLinearId::Next,  2, u32>;

            static constexpr u16 MinLinearId = 1;
            static constexpr u16 MaxLinearId = util::BitPack32{std::numeric_limits<u32>::max()}.Get<HandleLinearId>();

            static constexpr ALWAYS_INLINE ams::svc::Handle EncodeHandle(u16 index, u16 linear_id) {
                util::BitPack32 pack = {0};
                pack.Set<HandleIndex>(index);
                pack.Set<HandleLinearId>(linear_id);
                pack.Set<HandleReserved>(0);
                return pack.Get<HandleEncoded>();
            }

            class Entry {
                private:
                    union {
                        struct {
                            u16 linear_id;
                            u16 type;
                        } info;
                        Entry *next_free_entry;
                    } m_meta;
                    KAutoObject *m_object;
                public:
                    constexpr Entry() : m_meta(), m_object(nullptr) { /* ... */ }

                    constexpr ALWAYS_INLINE void SetFree(Entry *next) {
                        m_object = nullptr;
                        m_meta.next_free_entry = next;
                    }

                    constexpr ALWAYS_INLINE void SetUsed(KAutoObject *obj, u16 linear_id, u16 type) {
                        m_object = obj;
                        m_meta.info = { linear_id, type };
                    }

                    constexpr ALWAYS_INLINE KAutoObject *GetObject() const { return m_object; }
                    constexpr ALWAYS_INLINE Entry *GetNextFreeEntry() const { return m_meta.next_free_entry; }
                    constexpr ALWAYS_INLINE u16 GetLinearId() const { return m_meta.info.linear_id; }
                    constexpr ALWAYS_INLINE u16 GetType() const { return m_meta.info.type; }
            };
        private:
            mutable KSpinLock m_lock;
            Entry *m_table;
            Entry *m_free_head;
            Entry m_entries[MaxTableSize];
            u16 m_table_size;
            u16 m_max_count;
            u16 m_next_linear_id;
            u16 m_count;
        public:
            constexpr KHandleTable() :
                m_lock(), m_table(nullptr), m_free_head(nullptr), m_entries(), m_table_size(0), m_max_count(0), m_next_linear_id(MinLinearId), m_count(0)
            { MESOSPHERE_ASSERT_THIS(); }

            constexpr NOINLINE Result Initialize(s32 size) {
                MESOSPHERE_ASSERT_THIS();

                R_UNLESS(size <= static_cast<s32>(MaxTableSize), svc::ResultOutOfMemory());

                /* Initialize all fields. */
                m_table = m_entries;
                m_table_size = (size <= 0) ? MaxTableSize : size;
                m_next_linear_id = MinLinearId;
                m_count = 0;
                m_max_count = 0;

                /* Free all entries. */
                for (size_t i = 0; i < static_cast<size_t>(m_table_size - 1); i++) {
                    m_entries[i].SetFree(std::addressof(m_entries[i + 1]));
                }
                m_entries[m_table_size - 1].SetFree(nullptr);

                m_free_head = std::addressof(m_entries[0]);

                return ResultSuccess();
            }

            constexpr ALWAYS_INLINE size_t GetTableSize() const { return m_table_size; }
            constexpr ALWAYS_INLINE size_t GetCount() const { return m_count; }
            constexpr ALWAYS_INLINE size_t GetMaxCount() const { return m_max_count; }

            NOINLINE Result Finalize();
            NOINLINE bool Remove(ams::svc::Handle handle);

            template<typename T = KAutoObject>
            ALWAYS_INLINE KScopedAutoObject<T> GetObjectWithoutPseudoHandle(ams::svc::Handle handle) const {
                /* Lock and look up in table. */
                KScopedDisableDispatch dd;
                KScopedSpinLock lk(m_lock);

                if constexpr (std::is_same<T, KAutoObject>::value) {
                    return this->GetObjectImpl(handle);
                } else {
                    if (auto *obj = this->GetObjectImpl(handle); obj != nullptr) {
                        return obj->DynamicCast<T*>();
                    } else {
                        return nullptr;
                    }
                }
            }

            template<typename T = KAutoObject>
            ALWAYS_INLINE KScopedAutoObject<T> GetObject(ams::svc::Handle handle) const {
                MESOSPHERE_ASSERT_THIS();

                /* Handle pseudo-handles. */
                if constexpr (std::derived_from<KProcess, T>) {
                    if (handle == ams::svc::PseudoHandle::CurrentProcess) {
                        return GetCurrentProcessPointer();
                    }
                } else if constexpr (std::derived_from<KThread, T>) {
                    if (handle == ams::svc::PseudoHandle::CurrentThread) {
                        return GetCurrentThreadPointer();
                    }
                }

                return this->template GetObjectWithoutPseudoHandle<T>(handle);
            }

            KScopedAutoObject<KAutoObject> GetObjectForIpcWithoutPseudoHandle(ams::svc::Handle handle) const {
                /* Lock and look up in table. */
                KScopedDisableDispatch dd;
                KScopedSpinLock lk(m_lock);

                KAutoObject *obj = this->GetObjectImpl(handle);
                if (AMS_LIKELY(obj != nullptr)) {
                    if (AMS_UNLIKELY(obj->DynamicCast<KInterruptEvent *>() != nullptr)) {
                        return nullptr;
                    }
                }

                return obj;
            }

            ALWAYS_INLINE KScopedAutoObject<KAutoObject> GetObjectForIpc(ams::svc::Handle handle, KThread *cur_thread) const {
                /* Handle pseudo-handles. */
                if (handle == ams::svc::PseudoHandle::CurrentProcess) {
                    return static_cast<KAutoObject *>(static_cast<void *>(cur_thread->GetOwnerProcess()));
                }
                if (handle == ams::svc::PseudoHandle::CurrentThread) {
                    return static_cast<KAutoObject *>(cur_thread);
                }

                return GetObjectForIpcWithoutPseudoHandle(handle);
            }

            ALWAYS_INLINE KScopedAutoObject<KAutoObject> GetObjectByIndex(ams::svc::Handle *out_handle, size_t index) const {
                MESOSPHERE_ASSERT_THIS();
                KScopedDisableDispatch dd;
                KScopedSpinLock lk(m_lock);

                return this->GetObjectByIndexImpl(out_handle, index);
            }

            NOINLINE Result Reserve(ams::svc::Handle *out_handle);
            NOINLINE void Unreserve(ams::svc::Handle handle);

            template<typename T>
            ALWAYS_INLINE Result Add(ams::svc::Handle *out_handle, T *obj) {
                static_assert(std::is_base_of<KAutoObject, T>::value);
                return this->Add(out_handle, obj, obj->GetTypeObj().GetClassToken());
            }

            template<typename T>
            ALWAYS_INLINE void Register(ams::svc::Handle handle, T *obj) {
                static_assert(std::is_base_of<KAutoObject, T>::value);
                return this->Register(handle, obj, obj->GetTypeObj().GetClassToken());
            }

            template<typename T>
            ALWAYS_INLINE bool GetMultipleObjects(T **out, const ams::svc::Handle *handles, size_t num_handles) const {
                /* Try to convert and open all the handles. */
                size_t num_opened;
                {
                    /* Lock the table. */
                    KScopedDisableDispatch dd;
                    KScopedSpinLock lk(m_lock);
                    for (num_opened = 0; num_opened < num_handles; num_opened++) {
                        /* Get the current handle. */
                        const auto cur_handle = handles[num_opened];

                        /* Get the object for the current handle. */
                        KAutoObject *cur_object = this->GetObjectImpl(cur_handle);
                        if (AMS_UNLIKELY(cur_object == nullptr)) {
                            break;
                        }

                        /* Cast the current object to the desired type. */
                        T *cur_t = cur_object->DynamicCast<T*>();
                        if (AMS_UNLIKELY(cur_t == nullptr)) {
                            break;
                        }

                        /* Open a reference to the current object. */
                        cur_t->Open();
                        out[num_opened] = cur_t;
                    }
                }

                /* If we converted every object, succeed. */
                if (AMS_LIKELY(num_opened == num_handles)) {
                    return true;
                }

                /* If we didn't convert entry object, close the ones we opened. */
                for (size_t i = 0; i < num_opened; i++) {
                    out[i]->Close();
                }

                return false;
            }
        private:
            NOINLINE Result Add(ams::svc::Handle *out_handle, KAutoObject *obj, u16 type);
            NOINLINE void Register(ams::svc::Handle handle, KAutoObject *obj, u16 type);

            constexpr ALWAYS_INLINE Entry *AllocateEntry() {
                MESOSPHERE_ASSERT_THIS();
                MESOSPHERE_ASSERT(m_count < m_table_size);

                Entry *entry = m_free_head;
                m_free_head = entry->GetNextFreeEntry();

                m_count++;
                m_max_count = std::max(m_max_count, m_count);

                return entry;
            }

            constexpr ALWAYS_INLINE void FreeEntry(Entry *entry) {
                MESOSPHERE_ASSERT_THIS();
                MESOSPHERE_ASSERT(m_count > 0);

                entry->SetFree(m_free_head);
                m_free_head = entry;

                m_count--;
            }

            constexpr ALWAYS_INLINE u16 AllocateLinearId() {
                const u16 id = m_next_linear_id++;
                if (m_next_linear_id > MaxLinearId) {
                    m_next_linear_id = MinLinearId;
                }
                return id;
            }

            constexpr ALWAYS_INLINE size_t GetEntryIndex(Entry *entry) {
                const size_t index = entry - m_table;
                MESOSPHERE_ASSERT(index < m_table_size);
                return index;
            }

            constexpr ALWAYS_INLINE Entry *FindEntry(ams::svc::Handle handle) const {
                MESOSPHERE_ASSERT_THIS();

                /* Unpack the handle. */
                const auto handle_pack = GetHandleBitPack(handle);
                const auto raw_value   = handle_pack.Get<HandleRawValue>();
                const auto index       = handle_pack.Get<HandleIndex>();
                const auto linear_id   = handle_pack.Get<HandleLinearId>();
                const auto reserved    = handle_pack.Get<HandleReserved>();
                MESOSPHERE_ASSERT(reserved == 0);
                MESOSPHERE_UNUSED(reserved);

                /* Validate our indexing information. */
                if (raw_value == 0) {
                    return nullptr;
                }
                if (linear_id == 0) {
                    return nullptr;
                }
                if (index >= m_table_size) {
                    return nullptr;
                }

                /* Get the entry, and ensure our serial id is correct. */
                Entry *entry = std::addressof(m_table[index]);
                if (entry->GetObject() == nullptr) {
                    return nullptr;
                }
                if (entry->GetLinearId() != linear_id) {
                    return nullptr;
                }

                return entry;
            }

            constexpr ALWAYS_INLINE KAutoObject *GetObjectImpl(ams::svc::Handle handle) const {
                MESOSPHERE_ASSERT_THIS();

                /* Handles must not have reserved bits set. */
                if (GetHandleBitPack(handle).Get<HandleReserved>() != 0) {
                    return nullptr;
                }

                if (Entry *entry = this->FindEntry(handle); entry != nullptr) {
                    return entry->GetObject();
                } else {
                    return nullptr;
                }
            }

            constexpr ALWAYS_INLINE KAutoObject *GetObjectByIndexImpl(ams::svc::Handle *out_handle, size_t index) const {
                MESOSPHERE_ASSERT_THIS();

                /* Index must be in bounds. */
                if (index >= m_table_size || m_table == nullptr) {
                    return nullptr;
                }

                /* Ensure entry has an object. */
                Entry *entry = std::addressof(m_table[index]);
                if (entry->GetObject() == nullptr) {
                    return nullptr;
                }

                *out_handle = EncodeHandle(index, entry->GetLinearId());
                return entry->GetObject();
            }
    };

}