Kernel: Remove storage devices after unplugging them

kernel/include/kernel/Device/Device.h

@@ -28,7 +28,7 @@ namespace Kernel
 		virtual BAN::ErrorOr<void> fsync_impl() final override { return BAN::Error::from_errno(EINVAL); }
 	};
 
-	class BlockDevice : public Device
+	class BlockDevice : public Device, public BAN::Weakable<BlockDevice>
 	{
 	public:
 		virtual BAN::ErrorOr<void> read_blocks(uint64_t first_block, size_t block_count, BAN::ByteSpan) = 0;

kernel/include/kernel/Storage/Partition.h

@@ -13,9 +13,8 @@ namespace Kernel
 
 		const BAN::GUID& partition_type() const { return m_type; }
 		const BAN::GUID& partition_guid() const { return m_guid; }
-		const BAN::RefPtr<BlockDevice> device() const { return m_device; }
 
-		virtual blksize_t blksize() const override { return m_device->blksize(); }
+		virtual blksize_t blksize() const override { return m_block_size; }
 
 		BAN::ErrorOr<void> read_sectors(uint64_t first_block, size_t block_count, BAN::ByteSpan buffer)			{ return read_blocks(first_block, block_count, buffer); }
 		BAN::ErrorOr<void> write_sectors(uint64_t first_block, size_t block_count, BAN::ConstByteSpan buffer)	{ return write_blocks(first_block, block_count, buffer); }
@@ -32,7 +31,8 @@ namespace Kernel
 		Partition(BAN::RefPtr<BlockDevice>, const BAN::GUID&, const BAN::GUID&, uint64_t, uint64_t, uint64_t, const char*, uint32_t, BAN::StringView);
 
 	private:
-		BAN::RefPtr<BlockDevice> m_device;
+		mutable BAN::WeakPtr<BlockDevice> m_device;
+		const uint64_t m_block_size;
 		const BAN::GUID m_type;
 		const BAN::GUID m_guid;
 		const BAN::String m_guid_string;

kernel/kernel/Storage/Partition.cpp

@@ -16,6 +16,7 @@ namespace Kernel
 	Partition::Partition(BAN::RefPtr<BlockDevice> device, const BAN::GUID& type, const BAN::GUID& guid, uint64_t first_block, uint64_t last_block, uint64_t attr, const char* label, uint32_t index, BAN::StringView name_prefix)
 		: BlockDevice(0660, 0, 0)
 		, m_device(device)
+		, m_block_size(device->blksize())
 		, m_type(type)
 		, m_guid(guid)
 		, m_guid_string(MUST(guid.to_string()))
@@ -30,30 +31,42 @@ namespace Kernel
 
 	BAN::ErrorOr<void> Partition::read_blocks(uint64_t first_block, size_t block_count, BAN::ByteSpan buffer)
 	{
-		ASSERT(buffer.size() >= block_count * m_device->blksize());
+		auto device = m_device.lock();
+		if (!device)
+			return BAN::Error::from_errno(ENODEV);
+
+		ASSERT(buffer.size() >= block_count * device->blksize());
 		const uint32_t blocks_in_partition = m_last_block - m_first_block + 1;
 		if (first_block + block_count > blocks_in_partition)
 			return BAN::Error::from_error_code(ErrorCode::Storage_Boundaries);
-		TRY(m_device->read_blocks(m_first_block + first_block, block_count, buffer));
+		TRY(device->read_blocks(m_first_block + first_block, block_count, buffer));
 		return {};
 	}
 
 	BAN::ErrorOr<void> Partition::write_blocks(uint64_t first_block, size_t block_count, BAN::ConstByteSpan buffer)
 	{
-		ASSERT(buffer.size() >= block_count * m_device->blksize());
+		auto device = m_device.lock();
+		if (!device)
+			return BAN::Error::from_errno(ENODEV);
+
+		ASSERT(buffer.size() >= block_count * device->blksize());
 		const uint32_t blocks_in_partition = m_last_block - m_first_block + 1;
 		if (first_block + block_count > blocks_in_partition)
 			return BAN::Error::from_error_code(ErrorCode::Storage_Boundaries);
-		TRY(m_device->write_blocks(m_first_block + first_block, block_count, buffer));
+		TRY(device->write_blocks(m_first_block + first_block, block_count, buffer));
 		return {};
 	}
 
 	BAN::ErrorOr<void> Partition::sync_blocks(uint64_t block, size_t block_count)
 	{
+		auto device = m_device.lock();
+		if (!device)
+			return BAN::Error::from_errno(ENODEV);
+
 		const uint32_t blocks_in_partition = m_last_block - m_first_block + 1;
 		if (block + block_count > blocks_in_partition)
 			return BAN::Error::from_error_code(ErrorCode::Storage_Boundaries);
-		TRY(m_device->sync_blocks(m_first_block + block, block_count));
+		TRY(device->sync_blocks(m_first_block + block, block_count));
 		return {};
 	}
 
@@ -61,12 +74,16 @@ namespace Kernel
 	{
 		ASSERT(offset >= 0);
 
-		if (offset % m_device->blksize() || buffer.size() % m_device->blksize())
+		auto device = m_device.lock();
+		if (!device)
+			return BAN::Error::from_errno(ENODEV);
+
+		if (offset % device->blksize() || buffer.size() % device->blksize())
 			return BAN::Error::from_errno(ENOTSUP);
 
 		const uint32_t blocks_in_partition = m_last_block - m_first_block + 1;
-		uint32_t first_block = offset / m_device->blksize();
+		uint32_t first_block = offset / device->blksize();
-		uint32_t block_count = buffer.size() / m_device->blksize();
+		uint32_t block_count = buffer.size() / device->blksize();
 
 		if (first_block >= blocks_in_partition)
 			return 0;
@@ -74,7 +91,7 @@ namespace Kernel
 			block_count = blocks_in_partition - first_block;
 
 		TRY(read_blocks(first_block, block_count, buffer));
-		return block_count * m_device->blksize();
+		return block_count * device->blksize();
 	}
 
 }

kernel/kernel/Storage/StorageDevice.cpp

@@ -204,6 +204,9 @@ namespace Kernel
 
 	StorageDevice::~StorageDevice()
 	{
+		for (auto& partition : m_partitions)
+			if (partition)
+				DevFileSystem::get().remove_device(partition);
 	}
 
 	void StorageDevice::add_disk_cache()

kernel/kernel/USB/MassStorage/MassStorageDriver.cpp

@@ -20,7 +20,11 @@ namespace Kernel
 	{ }
 
 	USBMassStorageDriver::~USBMassStorageDriver()
-	{ }
+	{
+		for (auto& device : m_storage_devices)
+			if (device)
+				DevFileSystem::get().remove_device(device);
+	}
 
 	BAN::ErrorOr<void> USBMassStorageDriver::initialize()
 	{