From a1783740526c64cf576a8495aa65a11085af63a8 Mon Sep 17 00:00:00 2001
From: Eladash <elad3356p@gmail.com>
Date: Sat, 4 Apr 2020 17:11:21 +0300
Subject: [PATCH] sys_fs: Limit NPDRM FDs to 16

Co-Authored-By: Silent <cookieplmonster@users.noreply.github.com>
---
 rpcs3/Crypto/unedat.h         |  1 +
 rpcs3/Emu/Cell/lv2/sys_fs.cpp | 55 +++++++++++++++++++++++++++++++++--
 rpcs3/Emu/Cell/lv2/sys_fs.h   | 13 +++++++--
 3 files changed, 64 insertions(+), 5 deletions(-)

diff --git a/rpcs3/Crypto/unedat.h b/rpcs3/Crypto/unedat.h
index 0bde5e7104..94788e497f 100644
--- a/rpcs3/Crypto/unedat.h
+++ b/rpcs3/Crypto/unedat.h
@@ -18,6 +18,7 @@ struct loaded_npdrm_keys
 {
 	std::array<u8, 0x10> devKlic{};
 	std::array<u8, 0x10> rifKey{};
+	atomic_t<u32> npdrm_fds{0};
 };
 
 struct NPD_HEADER
diff --git a/rpcs3/Emu/Cell/lv2/sys_fs.cpp b/rpcs3/Emu/Cell/lv2/sys_fs.cpp
index d901ed7ec1..c93e5cd494 100644
--- a/rpcs3/Emu/Cell/lv2/sys_fs.cpp
+++ b/rpcs3/Emu/Cell/lv2/sys_fs.cpp
@@ -485,7 +485,42 @@ error_code sys_fs_open(ppu_thread& ppu, vm::cptr<char> path, s32 flags, vm::ptr<
 		return {error, path};
 	}
 
-	if (const u32 id = idm::make<lv2_fs_object, lv2_file>(ppath, std::move(file), mode, flags))
+	lv2_file_type type = lv2_file_type::regular;
+
+	if (size == 8)
+	{
+		// see lv2_file::open
+		switch (vm::read64(arg.addr()))
+		{
+		case 0x18000000010:
+		case 0x2:
+		{
+			type = lv2_file_type::npdrm;
+			sys_fs.warning("sys_fs_open(): NPDRM detected");
+			break;
+		}
+		default:
+			break;
+		}
+	}
+
+	if (type == lv2_file_type::npdrm)
+	{
+		if (const u32 id = idm::import<lv2_fs_object, lv2_file>([&ppath = ppath, &file = file, mode, flags]() -> std::shared_ptr<lv2_file>
+		{
+			if (!g_fxo->get<loaded_npdrm_keys>()->npdrm_fds.try_inc(16))
+			{
+				return nullptr;
+			}
+
+			return std::make_shared<lv2_file>(ppath, std::move(file), mode, flags, lv2_file_type::npdrm);
+		}))
+		{
+			*fd = id;
+			return CELL_OK;
+		}
+	}
+	else if (const u32 id = idm::make<lv2_fs_object, lv2_file>(ppath, std::move(file), mode, flags))
 	{
 		*fd = id;
 		return CELL_OK;
@@ -594,7 +629,13 @@ error_code sys_fs_close(ppu_thread& ppu, u32 fd)
 
 	sys_fs.trace("sys_fs_close(fd=%d)", fd);
 
-	const auto file = idm::withdraw<lv2_fs_object, lv2_file>(fd);
+	const auto file = idm::withdraw<lv2_fs_object, lv2_file>(fd, [](lv2_file& file)
+	{
+		if (file.type == lv2_file_type::npdrm)
+		{
+			g_fxo->get<loaded_npdrm_keys>()->npdrm_fds--;
+		}
+	});
 
 	if (!file)
 	{
@@ -1236,7 +1277,15 @@ error_code sys_fs_fcntl(ppu_thread& ppu, u32 fd, u32 op, vm::ptr<void> _arg, u32
 
 		fs::file stream;
 		stream.reset(std::move(sdata_file));
-		if (const u32 id = idm::make<lv2_fs_object, lv2_file>(*file, std::move(stream), file->mode, file->flags))
+		if (const u32 id = idm::import<lv2_fs_object, lv2_file>([&]() -> std::shared_ptr<lv2_file>
+		{
+			if (!g_fxo->get<loaded_npdrm_keys>()->npdrm_fds.try_inc(16))
+			{
+				return nullptr;
+			}
+
+			return std::make_shared<lv2_file>(*file, std::move(stream), file->mode, file->flags, lv2_file_type::npdrm);
+		}))
 		{
 			arg->out_code = CELL_OK;
 			arg->out_fd = id;
diff --git a/rpcs3/Emu/Cell/lv2/sys_fs.h b/rpcs3/Emu/Cell/lv2/sys_fs.h
index d62cd39eda..e78eebac26 100644
--- a/rpcs3/Emu/Cell/lv2/sys_fs.h
+++ b/rpcs3/Emu/Cell/lv2/sys_fs.h
@@ -132,6 +132,12 @@ enum class lv2_mp_flag
 	__bitset_enum_max
 };
 
+enum class lv2_file_type
+{
+	regular = 0,
+	npdrm,
+};
+
 struct lv2_fs_object
 {
 	using id_type = lv2_fs_object;
@@ -174,23 +180,26 @@ struct lv2_file final : lv2_fs_object
 	const fs::file file;
 	const s32 mode;
 	const s32 flags;
+	const lv2_file_type type;
 
 	// Stream lock
 	atomic_t<u32> lock{0};
 
-	lv2_file(std::string_view filename, fs::file&& file, s32 mode, s32 flags)
+	lv2_file(std::string_view filename, fs::file&& file, s32 mode, s32 flags, lv2_file_type type = {})
 		: lv2_fs_object(lv2_fs_object::get_mp(filename), filename)
 		, file(std::move(file))
 		, mode(mode)
 		, flags(flags)
+		, type(type)
 	{
 	}
 
-	lv2_file(const lv2_file& host, fs::file&& file, s32 mode, s32 flags)
+	lv2_file(const lv2_file& host, fs::file&& file, s32 mode, s32 flags, lv2_file_type type = {})
 		: lv2_fs_object(host.mp, host.name.data())
 		, file(std::move(file))
 		, mode(mode)
 		, flags(flags)
+		, type(type)
 	{
 	}