orbis: initial budget implementation

orbis-kernel/include/orbis/AuthInfo.hpp

@@ -9,9 +9,13 @@ struct AuthInfo {
   uint64_t attrs[4];
   uint64_t ucred[8];
 
+  bool isSyscoreProcess() const { return ucred[2] == 0x3800000000000007; }
+  bool isShellUiProcess() const { return ucred[2] == 0x380000000000000f; }
+
   bool hasUseHp3dPipeCapability() const {
     return ucred[2] == 0x3800000000000009;
   }
+
   bool hasMmapSelfCapability() const { return ((ucred[4] >> 0x3a) & 1) != 1; }
   bool hasSystemCapability() const { return ((ucred[3] >> 0x3e) & 1) != 0; }
   bool hasSceProgramAttribute() const { return ((ucred[3] >> 0x1f) & 1) != 0; }

orbis-kernel/include/orbis/Budget.hpp

@@ -0,0 +1,141 @@
+#pragma once
+
+#include "orbis-config.hpp"
+#include "utils/BitSet.hpp"
+#include "utils/Rc.hpp"
+#include "utils/SharedMutex.hpp"
+#include <array>
+#include <cstring>
+#include <mutex>
+#include <string_view>
+
+namespace orbis {
+enum class BudgetResource : uint32_t {
+  Invalid,
+  Dmem,
+  Vmem,
+  Fmem,
+  CpuSet,
+  File,
+  Socket,
+  Equeue,
+  Pipe,
+  Device,
+  Thread,
+  IpSocket,
+
+  _count,
+};
+
+struct BudgetItem {
+  uint64_t total;
+  uint64_t used;
+};
+
+struct BudgetInfo {
+  BudgetResource resourceId;
+  uint32_t flags; // ?
+  BudgetItem item;
+};
+
+static_assert(sizeof(BudgetInfo) == 0x18);
+
+using BudgetInfoList =
+    std::array<BudgetInfo, static_cast<int>(BudgetResource::_count)>;
+
+class Budget : public RcBase {
+  using BudgetList =
+      std::array<BudgetItem, static_cast<int>(BudgetResource::_count)>;
+
+public:
+  enum class ProcessType : orbis::uint32_t {
+    BigApp,
+    MiniApp,
+    System,
+    NonGameMiniApp,
+    _last = NonGameMiniApp
+  };
+
+  Budget(std::string_view name, ProcessType pType,
+         std::span<const BudgetInfo> budgets)
+      : mProcessType(pType) {
+    for (auto info : budgets) {
+      if (info.resourceId == BudgetResource::Invalid) {
+        continue;
+      }
+
+      int resourceIndex = static_cast<int>(info.resourceId);
+
+      mUsed.set(resourceIndex);
+      mList[resourceIndex] = info.item;
+    }
+
+    std::strncpy(mName, name.data(), std::min(name.size(), sizeof(mName)));
+    mName[sizeof(mName) - 1] = 0;
+  }
+
+  [[nodiscard]] BudgetItem get(BudgetResource resourceId) const {
+    std::lock_guard lock(mMtx);
+    return mList[static_cast<int>(resourceId)];
+  }
+
+  [[nodiscard]] BudgetList getBudgetList() const {
+    std::lock_guard lock(mMtx);
+    return mList;
+  }
+
+  [[nodiscard]] std::pair<BudgetInfoList, int> getList() const {
+    auto budgetList = getBudgetList();
+
+    BudgetInfoList result{};
+    int count = 0;
+
+    for (auto resourceId : mUsed) {
+      result[count].resourceId = static_cast<BudgetResource>(resourceId);
+      result[count].item = budgetList[resourceId];
+      count++;
+    }
+
+    return {result, count};
+  }
+
+  bool acquire(BudgetResource resourceId, std::uint64_t size = 1) {
+    auto &budget = mList[static_cast<int>(resourceId)];
+
+    std::lock_guard lock(mMtx);
+
+    if (budget.used + size > budget.total) {
+      return false;
+    }
+
+    budget.used += size;
+    return true;
+  }
+
+  void release(BudgetResource resourceId, std::uint64_t size) {
+    auto &budget = mList[static_cast<int>(resourceId)];
+
+    std::lock_guard lock(mMtx);
+
+    if (size >= budget.used) {
+      budget.used = 0;
+    } else {
+      budget.used -= size;
+    }
+  }
+
+  bool hasResource(BudgetResource resourceId) const {
+    return mUsed.test(static_cast<int>(resourceId));
+  }
+
+  [[nodiscard]] int size() const { return mUsed.popcount(); }
+  [[nodiscard]] ProcessType processType() const { return mProcessType; }
+
+private:
+  mutable shared_mutex mMtx;
+  orbis::BitSet<static_cast<int>(BudgetResource::_count)> mUsed;
+  ProcessType mProcessType{};
+  BudgetList mList;
+  char mName[32]{};
+};
+} // namespace orbis

orbis-kernel/include/orbis/sys/sysproto.hpp

@@ -1,4 +1,5 @@
 #include "orbis-config.hpp"
+#include <orbis/Budget.hpp>
 #include <orbis/error.hpp>
 #include <orbis/module/ModuleHandle.hpp>
 #include <orbis/thread/cpuset.hpp>
@@ -669,11 +670,14 @@ SysResult sys_opmc_disable(Thread *thread /* TODO */);
 SysResult sys_opmc_set_ctl(Thread *thread /* TODO */);
 SysResult sys_opmc_set_ctr(Thread *thread /* TODO */);
 SysResult sys_opmc_get_ctr(Thread *thread /* TODO */);
-SysResult sys_budget_create(Thread *thread /* TODO */);
-SysResult sys_budget_delete(Thread *thread /* TODO */);
-SysResult sys_budget_get(Thread *thread, sint id, ptr<void> a,
-                         ptr<uint32_t> count);
-SysResult sys_budget_set(Thread *thread, slong budget);
+SysResult sys_budget_create(Thread *thread, ptr<char> name,
+                            Budget::ProcessType processType,
+                            ptr<const BudgetInfo> resources, orbis::uint count,
+                            ptr<BudgetInfo> invalidResources);
+SysResult sys_budget_delete(Thread *thread, sint budget);
+SysResult sys_budget_get(Thread *thread, sint id, ptr<BudgetInfo> budgetInfo,
+                         ptr<sint> count);
+SysResult sys_budget_set(Thread *thread, sint budget);
 SysResult sys_virtual_query(Thread *thread, ptr<void> addr, uint64_t unk,
                             ptr<void> info, size_t infosz);
 SysResult sys_mdbg_call(Thread *thread /* TODO */);
@@ -726,7 +730,7 @@ SysResult sys_dynlib_get_info_ex(Thread *thread, SceKernelModule handle,
                                  ptr<struct Unk> unk,
                                  ptr<ModuleInfoEx> destModuleInfoEx);
 SysResult sys_budget_getid(Thread *thread);
-SysResult sys_budget_get_ptype(Thread *thread, sint budgetId);
+SysResult sys_budget_get_ptype(Thread *thread, sint pid);
 SysResult sys_get_paging_stats_of_all_threads(Thread *thread /* TODO */);
 SysResult sys_get_proc_type_info(Thread *thread, ptr<sint> destProcessInfo);
 SysResult sys_get_resident_count(Thread *thread, pid_t pid);
@@ -773,7 +777,7 @@ SysResult sys_extend_page_table_pool(Thread *thread);
 SysResult sys_extend_page_table_pool2(Thread *thread);
 SysResult sys_get_kernel_mem_statistics(Thread *thread /* TODO */);
 SysResult sys_get_sdk_compiled_version(Thread *thread, ptr<const char> path);
-SysResult sys_app_state_change(Thread *thread /* TODO */);
+SysResult sys_app_state_change(Thread *thread, sint state);
 SysResult sys_dynlib_get_obj_member(Thread *thread, SceKernelModule handle,
                                     uint64_t index, ptr<ptr<void>> addrp);
 SysResult sys_budget_get_ptype_of_budget(Thread *thread, sint budgetId);

orbis-kernel/include/orbis/thread/Process.hpp

@@ -11,6 +11,7 @@
 #include "cpuset.hpp"
 #include "orbis/AppInfo.hpp"
 #include "orbis/AuthInfo.hpp"
+#include "orbis/Budget.hpp"
 #include "orbis/file.hpp"
 #include "orbis/module/Module.hpp"
 #include "orbis/utils/IdMap.hpp"
@@ -74,7 +75,8 @@ struct Process final {
   kstring root = "/";
   cpuset affinity{(1 << 7) - 1};
   sint memoryContainer{1};
-  sint budgetId{1};
+  sint budgetId{};
+  Budget::ProcessType budgetProcessType{};
   bool isInSandbox = false;
   EventEmitter event;
   std::optional<sint> exitStatus;

orbis-kernel/include/orbis/utils/BitSet.hpp

@@ -11,7 +11,31 @@ template <std::size_t Count> struct BitSet {
   static constexpr auto ChunkCount = (Count + BitsPerChunk - 1) / BitsPerChunk;
   chunk_type _bits[ChunkCount]{};
 
-  constexpr std::size_t countr_one() const {
+  struct iterator_end {};
+
+  struct iterator {
+    iterator() = default;
+
+    constexpr iterator &operator++() {
+      offset = bitSet->countr_zero(offset + 1);
+      return *this;
+    }
+    constexpr bool operator==(const iterator_end &) const {
+      return offset >= Count;
+    }
+    constexpr std::size_t operator*() const { return offset; }
+
+  private:
+    constexpr iterator(const BitSet *bitSet)
+        : bitSet(bitSet), offset(bitSet->countr_zero()) {}
+
+    const BitSet *bitSet = nullptr;
+    std::size_t offset = 0;
+
+    friend BitSet;
+  };
+
+  [[nodiscard]] constexpr std::size_t countr_one() const {
     std::size_t result = 0;
     for (auto bits : _bits) {
       auto count = std::countr_one(bits);
@@ -25,7 +49,18 @@ template <std::size_t Count> struct BitSet {
     return result;
   }
 
-  constexpr std::size_t countr_zero(std::size_t offset = 0) const {
+  [[nodiscard]] constexpr std::size_t popcount() const {
+    std::size_t result = 0;
+
+    for (auto bits : _bits) {
+      result += std::popcount(bits);
+    }
+
+    return result;
+  }
+
+  [[nodiscard]] constexpr std::size_t
+  countr_zero(std::size_t offset = 0) const {
     std::size_t result = 0;
 
     if (auto chunkOffset = offset % BitsPerChunk) {
@@ -47,21 +82,11 @@ template <std::size_t Count> struct BitSet {
         break;
       }
     }
-    /*
-    for (auto bits : _bits) {
-      auto count = std::countr_zero(bits);
-      result += count;
-
-      if (count != BitsPerChunk) {
-        break;
-      }
-    }
-    */
 
     return result + offset;
   }
 
-  bool empty() const {
+  [[nodiscard]] constexpr bool empty() const {
     for (auto bits : _bits) {
       if (bits != 0) {
         return false;
@@ -71,7 +96,7 @@ template <std::size_t Count> struct BitSet {
     return true;
   }
 
-  bool full() const {
+  [[nodiscard]] constexpr bool full() const {
     if constexpr (Count < BitsPerChunk) {
       return _bits[0] == (static_cast<std::uint64_t>(1) << Count) - 1;
     }
@@ -95,10 +120,13 @@ template <std::size_t Count> struct BitSet {
                                    << (index % BitsPerChunk);
   }
 
-  constexpr bool test(std::size_t index) const {
+  [[nodiscard]] constexpr bool test(std::size_t index) const {
     return (_bits[index / BitsPerChunk] &
             (static_cast<chunk_type>(1) << (index % BitsPerChunk))) != 0;
   }
+
+  [[nodiscard]] constexpr iterator begin() const { return iterator(this); }
+  [[nodiscard]] constexpr iterator_end end() const { return {}; }
 };
 } // namespace utils
 } // namespace orbis