From 1838ea5c306a690ff9e9a89abb3c4d10229db0e2 Mon Sep 17 00:00:00 2001
From: Bananymous <oskari.alaranta@bananymous.com>
Date: Mon, 4 Nov 2024 15:35:07 +0200
Subject: [PATCH] Kernel: Fix pseudo terminal leaks

---
 .../include/kernel/Terminal/PseudoTerminal.h  |  6 ++-
 kernel/kernel/Process.cpp                     |  4 +-
 kernel/kernel/Terminal/PseudoTerminal.cpp     | 44 ++++++++++++++++---
 3 files changed, 45 insertions(+), 9 deletions(-)

diff --git a/kernel/include/kernel/Terminal/PseudoTerminal.h b/kernel/include/kernel/Terminal/PseudoTerminal.h
index 26b0446b..1eebefb3 100644
--- a/kernel/include/kernel/Terminal/PseudoTerminal.h
+++ b/kernel/include/kernel/Terminal/PseudoTerminal.h
@@ -62,10 +62,12 @@ namespace Kernel
 		void putchar_impl(uint8_t ch) override;
 
 	private:
-		PseudoTerminalSlave(BAN::String&& name, mode_t, uid_t, gid_t);
+		PseudoTerminalSlave(BAN::String&& name, uint32_t number, mode_t, uid_t, gid_t);
+		~PseudoTerminalSlave();
 
 	private:
-		BAN::String m_name;
+		const BAN::String m_name;
+		const uint32_t m_number;
 
 		BAN::WeakPtr<PseudoTerminalMaster> m_master;
 		uint32_t m_width { 0 };
diff --git a/kernel/kernel/Process.cpp b/kernel/kernel/Process.cpp
index 1a0f9f6b..7a917f03 100644
--- a/kernel/kernel/Process.cpp
+++ b/kernel/kernel/Process.cpp
@@ -1790,13 +1790,13 @@ namespace Kernel
 		LockGuard _(m_process_lock);
 		TRY(validate_pointer_access(buffer, buffer_len, true));
 
-		auto inode = TRY(m_open_file_descriptors.inode_of(fildes));
 		if (TRY(m_open_file_descriptors.path_of(fildes)) != "<ptmx>"_sv)
 			return BAN::Error::from_errno(ENOTTY);
 
+		auto inode = TRY(m_open_file_descriptors.inode_of(fildes));
 		auto ptsname = TRY(static_cast<PseudoTerminalMaster*>(inode.ptr())->ptsname());
 
-		const size_t to_copy = BAN::Math::min(ptsname.size() + 1, buffer_len);
+		const size_t to_copy = BAN::Math::min(ptsname.size(), buffer_len - 1);
 		memcpy(buffer, ptsname.data(), to_copy);
 		buffer[to_copy] = '\0';
 
diff --git a/kernel/kernel/Terminal/PseudoTerminal.cpp b/kernel/kernel/Terminal/PseudoTerminal.cpp
index 73c1a0d5..51b4b5d7 100644
--- a/kernel/kernel/Terminal/PseudoTerminal.cpp
+++ b/kernel/kernel/Terminal/PseudoTerminal.cpp
@@ -2,16 +2,41 @@
 #include <kernel/FS/DevFS/FileSystem.h>
 #include <kernel/Terminal/PseudoTerminal.h>
 
+#include <BAN/ScopeGuard.h>
+
 #include <sys/sysmacros.h>
 
 namespace Kernel
 {
 
-	BAN::Atomic<uint32_t> s_pts_master_minor = 0;
-	BAN::Atomic<uint32_t> s_pts_slave_number = 0;
+	static BAN::Atomic<uint32_t> s_pts_master_minor = 0;
+
+	static SpinLock s_pts_slave_bitmap_lock;
+	static uint64_t s_pts_slave_bitmap = 0;
 
 	BAN::ErrorOr<BAN::RefPtr<PseudoTerminalMaster>> PseudoTerminalMaster::create(mode_t mode, uid_t uid, gid_t gid)
 	{
+		size_t pts_slave_number = 0;
+		{
+			SpinLockGuard _(s_pts_slave_bitmap_lock);
+			for (pts_slave_number = 0; pts_slave_number < sizeof(s_pts_slave_bitmap) * 8; pts_slave_number++)
+			{
+				if (s_pts_slave_bitmap & ((uint64_t)1 << pts_slave_number))
+					continue;
+				s_pts_slave_bitmap |= (uint64_t)1 << pts_slave_number;
+				break;
+			}
+		}
+		if (pts_slave_number >= sizeof(s_pts_slave_bitmap) * 8)
+			return BAN::Error::from_errno(ENODEV);
+		BAN::ScopeGuard slave_number_clearer(
+			[pts_slave_number]()
+			{
+				SpinLockGuard _(s_pts_slave_bitmap_lock);
+				s_pts_slave_bitmap &= ~((uint64_t)1 << pts_slave_number);
+			}
+		);
+
 		auto pts_master_buffer = TRY(VirtualRange::create_to_vaddr_range(
 			PageTable::kernel(),
 			KERNEL_OFFSET, static_cast<vaddr_t>(-1),
@@ -19,9 +44,11 @@ namespace Kernel
 			PageTable::Flags::ReadWrite | PageTable::Flags::Present, true
 		));
 		auto pts_master = TRY(BAN::RefPtr<PseudoTerminalMaster>::create(BAN::move(pts_master_buffer), mode, uid, gid));
+		DevFileSystem::get().remove_from_cache(pts_master);
 
-		auto pts_slave_name = TRY(BAN::String::formatted("pts{}", s_pts_slave_number++));
-		auto pts_slave = TRY(BAN::RefPtr<PseudoTerminalSlave>::create(BAN::move(pts_slave_name), 0610, uid, gid));
+		auto pts_slave_name = TRY(BAN::String::formatted("pts{}", pts_slave_number));
+		auto pts_slave = TRY(BAN::RefPtr<PseudoTerminalSlave>::create(BAN::move(pts_slave_name), pts_slave_number, 0610, uid, gid));
+		slave_number_clearer.disable();
 
 		pts_master->m_slave = TRY(pts_slave->get_weak_ptr());
 		pts_slave->m_master = TRY(pts_master->get_weak_ptr());
@@ -112,11 +139,18 @@ namespace Kernel
 		return buffer.size();
 	}
 
-	PseudoTerminalSlave::PseudoTerminalSlave(BAN::String&& name, mode_t mode, uid_t uid, gid_t gid)
+	PseudoTerminalSlave::PseudoTerminalSlave(BAN::String&& name, uint32_t number, mode_t mode, uid_t uid, gid_t gid)
 		: TTY(mode, uid, gid)
 		, m_name(BAN::move(name))
+		, m_number(number)
 	{}
 
+	PseudoTerminalSlave::~PseudoTerminalSlave()
+	{
+		SpinLockGuard _(s_pts_slave_bitmap_lock);
+		s_pts_slave_bitmap &= ~((uint64_t)1 << m_number);
+	}
+
 	void PseudoTerminalSlave::clear()
 	{
 		const char message[] { '\e', '[', '2', 'J' };