1
0
Fork 0
mirror of https://github.com/Atmosphere-NX/Atmosphere.git synced 2024-11-22 20:06:40 +00:00

kern: implement additional randomness in KPageHeap allocations

This commit is contained in:
Michael Scire 2022-03-22 15:29:55 -07:00 committed by SciresM
parent 24739f245e
commit 401047f603
6 changed files with 150 additions and 18 deletions

View file

@ -73,6 +73,7 @@ namespace ams::kern {
size_t Initialize(KPhysicalAddress address, size_t size, KVirtualAddress management, KVirtualAddress management_end, Pool p);
KPhysicalAddress AllocateBlock(s32 index, bool random) { return m_heap.AllocateBlock(index, random); }
KPhysicalAddress AllocateAligned(s32 index, size_t num_pages, size_t align_pages) { return m_heap.AllocateAligned(index, num_pages, align_pages); }
void Free(KPhysicalAddress addr, size_t num_pages) { m_heap.Free(addr, num_pages); }
void SetInitialUsedHeapSize(size_t reserved_size) { m_heap.SetInitialUsedSize(reserved_size); }

View file

@ -20,7 +20,7 @@
namespace ams::kern {
class KPageBitmap {
private:
public:
class RandomBitGenerator {
private:
util::TinyMT m_rng;
@ -42,12 +42,43 @@ namespace ams::kern {
--m_bits_available;
return rnd_bit;
}
u64 GenerateRandomBits(u32 num_bits) {
u64 result = 0;
/* Iteratively add random bits to our result. */
while (num_bits > 0) {
/* Ensure we have random bits to take from. */
if (m_bits_available == 0) {
this->RefreshEntropy();
}
/* Determine how many bits to take this round. */
const auto cur_bits = std::min(num_bits, m_bits_available);
/* Generate mask for our current bits. */
const u64 mask = (static_cast<u64>(1) << cur_bits) - 1;
/* Add bits to output from our entropy. */
result <<= cur_bits;
result |= (m_entropy & mask);
/* Remove bits from our entropy. */
m_entropy >>= cur_bits;
m_bits_available -= cur_bits;
/* Advance. */
num_bits -= cur_bits;
}
return result;
}
public:
RandomBitGenerator() : m_entropy(), m_bits_available() {
m_rng.Initialize(static_cast<u32>(KSystemControl::GenerateRandomU64()));
}
size_t SelectRandomBit(u64 bitmap) {
u64 SelectRandomBit(u64 bitmap) {
u64 selected = 0;
for (size_t cur_num_bits = BITSIZEOF(bitmap) / 2; cur_num_bits != 0; cur_num_bits /= 2) {
@ -66,6 +97,17 @@ namespace ams::kern {
return selected;
}
u64 GenerateRandom(u64 max) {
/* Determine the number of bits we need. */
const u64 bits_needed = 1 + (BITSIZEOF(max) - util::CountLeadingZeros(max));
/* Generate a random value of the desired bitwidth. */
const u64 rnd = this->GenerateRandomBits(bits_needed);
/* Adjust the value to be in range. */
return rnd - ((rnd / max) * max);
}
};
public:
static constexpr size_t MaxDepth = 4;

View file

@ -27,7 +27,7 @@ namespace ams::kern {
static constexpr s32 GetAlignedBlockIndex(size_t num_pages, size_t align_pages) {
const size_t target_pages = std::max(num_pages, align_pages);
for (size_t i = 0; i < NumMemoryBlockPageShifts; i++) {
if (target_pages <= (size_t(1) << MemoryBlockPageShifts[i]) / PageSize) {
if (target_pages <= (static_cast<size_t>(1) << MemoryBlockPageShifts[i]) / PageSize) {
return static_cast<s32>(i);
}
}
@ -36,7 +36,7 @@ namespace ams::kern {
static constexpr s32 GetBlockIndex(size_t num_pages) {
for (s32 i = static_cast<s32>(NumMemoryBlockPageShifts) - 1; i >= 0; i--) {
if (num_pages >= (size_t(1) << MemoryBlockPageShifts[i]) / PageSize) {
if (num_pages >= (static_cast<size_t>(1) << MemoryBlockPageShifts[i]) / PageSize) {
return i;
}
}
@ -44,7 +44,7 @@ namespace ams::kern {
}
static constexpr size_t GetBlockSize(size_t index) {
return size_t(1) << MemoryBlockPageShifts[index];
return static_cast<size_t>(1) << MemoryBlockPageShifts[index];
}
static constexpr size_t GetBlockNumPages(size_t index) {
@ -128,13 +128,14 @@ namespace ams::kern {
size_t m_initial_used_size;
size_t m_num_blocks;
Block m_blocks[NumMemoryBlockPageShifts];
KPageBitmap::RandomBitGenerator m_rng;
private:
void Initialize(KPhysicalAddress heap_address, size_t heap_size, KVirtualAddress management_address, size_t management_size, const size_t *block_shifts, size_t num_block_shifts);
size_t GetNumFreePages() const;
void FreeBlock(KPhysicalAddress block, s32 index);
public:
KPageHeap() : m_heap_address(Null<KPhysicalAddress>), m_heap_size(), m_initial_used_size(), m_num_blocks(), m_blocks() { /* ... */ }
KPageHeap() : m_heap_address(Null<KPhysicalAddress>), m_heap_size(), m_initial_used_size(), m_num_blocks(), m_blocks(), m_rng() { /* ... */ }
constexpr KPhysicalAddress GetAddress() const { return m_heap_address; }
constexpr size_t GetSize() const { return m_heap_size; }
@ -158,9 +159,25 @@ namespace ams::kern {
m_initial_used_size = m_heap_size - free_size - reserved_size;
}
KPhysicalAddress AllocateBlock(s32 index, bool random);
KPhysicalAddress AllocateBlock(s32 index, bool random) {
if (random) {
const size_t block_pages = m_blocks[index].GetNumPages();
return this->AllocateByRandom(index, block_pages, block_pages);
} else {
return this->AllocateByLinearSearch(index);
}
}
KPhysicalAddress AllocateAligned(s32 index, size_t num_pages, size_t align_pages) {
/* TODO: linear search support? */
return this->AllocateByRandom(index, num_pages, align_pages);
}
void Free(KPhysicalAddress addr, size_t num_pages);
private:
KPhysicalAddress AllocateByLinearSearch(s32 index);
KPhysicalAddress AllocateByRandom(s32 index, size_t num_pages, size_t align_pages);
static size_t CalculateManagementOverheadSize(size_t region_size, const size_t *block_shifts, size_t num_block_shifts);
public:
static size_t CalculateManagementOverheadSize(size_t region_size) {

View file

@ -202,7 +202,7 @@ namespace ams::kern {
Impl *chosen_manager = nullptr;
KPhysicalAddress allocated_block = Null<KPhysicalAddress>;
for (chosen_manager = this->GetFirstManager(pool, dir); chosen_manager != nullptr; chosen_manager = this->GetNextManager(chosen_manager, dir)) {
allocated_block = chosen_manager->AllocateBlock(heap_index, true);
allocated_block = chosen_manager->AllocateAligned(heap_index, num_pages, align_pages);
if (allocated_block != Null<KPhysicalAddress>) {
break;
}
@ -213,12 +213,6 @@ namespace ams::kern {
return Null<KPhysicalAddress>;
}
/* If we allocated more than we need, free some. */
const size_t allocated_pages = KPageHeap::GetBlockNumPages(heap_index);
if (allocated_pages > num_pages) {
chosen_manager->Free(allocated_block + num_pages * PageSize, allocated_pages - num_pages);
}
/* Maintain the optimized memory bitmap, if we should. */
if (m_has_optimized_process[pool]) {
chosen_manager->TrackUnoptimizedAllocation(allocated_block, num_pages);

View file

@ -51,11 +51,11 @@ namespace ams::kern {
return num_free;
}
KPhysicalAddress KPageHeap::AllocateBlock(s32 index, bool random) {
KPhysicalAddress KPageHeap::AllocateByLinearSearch(s32 index) {
const size_t needed_size = m_blocks[index].GetSize();
for (s32 i = index; i < static_cast<s32>(m_num_blocks); i++) {
if (const KPhysicalAddress addr = m_blocks[i].PopBlock(random); addr != Null<KPhysicalAddress>) {
if (const KPhysicalAddress addr = m_blocks[i].PopBlock(false); addr != Null<KPhysicalAddress>) {
if (const size_t allocated_size = m_blocks[i].GetSize(); allocated_size > needed_size) {
this->Free(addr + needed_size, (allocated_size - needed_size) / PageSize);
}
@ -66,6 +66,84 @@ namespace ams::kern {
return Null<KPhysicalAddress>;
}
KPhysicalAddress KPageHeap::AllocateByRandom(s32 index, size_t num_pages, size_t align_pages) {
/* Get the size and required alignment. */
const size_t needed_size = num_pages * PageSize;
const size_t align_size = align_pages * PageSize;
/* Determine meta-alignment of our desired alignment size. */
const size_t align_shift = util::CountTrailingZeros(align_size);
/* Decide on a block to allocate from. */
constexpr size_t MinimumPossibleAlignmentsForRandomAllocation = 4;
{
/* By default, we'll want to look at all blocks larger than our current one. */
s32 max_blocks = static_cast<s32>(m_num_blocks);
/* Determine the maximum block we should try to allocate from. */
size_t possible_alignments = 0;
for (s32 i = index; i < max_blocks; ++i) {
/* Add the possible alignments from blocks at the current size. */
possible_alignments += (1 + ((m_blocks[i].GetSize() - needed_size) >> align_shift)) * m_blocks[i].GetNumFreeBlocks();
/* If there are enough possible alignments, we don't need to look at larger blocks. */
if (possible_alignments >= MinimumPossibleAlignmentsForRandomAllocation) {
max_blocks = i + 1;
break;
}
}
/* If we have any possible alignments which require a larger block, we need to pick one. */
if (possible_alignments > 0 && index + 1 < max_blocks) {
/* Select a random alignment from the possibilities. */
const size_t rnd = m_rng.GenerateRandom(possible_alignments);
/* Determine which block corresponds to the random alignment we chose. */
possible_alignments = 0;
for (s32 i = index; i < max_blocks; ++i) {
/* Add the possible alignments from blocks at the current size. */
possible_alignments += (1 + ((m_blocks[i].GetSize() - needed_size) >> align_shift)) * m_blocks[i].GetNumFreeBlocks();
/* If the current block gets us to our random choice, use the current block. */
if (rnd < possible_alignments) {
index = i;
break;
}
}
}
}
/* Pop a block from the index we selected. */
if (KPhysicalAddress addr = m_blocks[index].PopBlock(true); addr != Null<KPhysicalAddress>) {
/* Determine how much size we have left over. */
if (const size_t leftover_size = m_blocks[index].GetSize() - needed_size; leftover_size > 0) {
/* Determine how many valid alignments we can have. */
const size_t possible_alignments = 1 + (leftover_size >> align_shift);
/* Select a random valid alignment. */
const size_t random_offset = m_rng.GenerateRandom(possible_alignments) << align_shift;
/* Free memory before the random offset. */
if (random_offset != 0) {
this->Free(addr, random_offset / PageSize);
}
/* Advance our block by the random offset. */
addr += random_offset;
/* Free memory after our allocated block. */
if (random_offset != leftover_size) {
this->Free(addr + needed_size, (leftover_size - random_offset) / PageSize);
}
}
/* Return the block we allocated. */
return addr;
}
return Null<KPhysicalAddress>;
}
void KPageHeap::FreeBlock(KPhysicalAddress block, s32 index) {
do {
block = m_blocks[index++].PushBlock(block);

View file

@ -3608,13 +3608,13 @@ namespace ams::kern {
/* Allocate the start page as needed. */
if (aligned_src_start < mapping_src_start) {
start_partial_page = Kernel::GetMemoryManager().AllocateAndOpenContinuous(1, 0, m_allocate_option);
start_partial_page = Kernel::GetMemoryManager().AllocateAndOpenContinuous(1, 1, m_allocate_option);
R_UNLESS(start_partial_page != Null<KPhysicalAddress>, svc::ResultOutOfMemory());
}
/* Allocate the end page as needed. */
if (mapping_src_end < aligned_src_end && (aligned_src_start < mapping_src_end || aligned_src_start == mapping_src_start)) {
end_partial_page = Kernel::GetMemoryManager().AllocateAndOpenContinuous(1, 0, m_allocate_option);
end_partial_page = Kernel::GetMemoryManager().AllocateAndOpenContinuous(1, 1, m_allocate_option);
R_UNLESS(end_partial_page != Null<KPhysicalAddress>, svc::ResultOutOfMemory());
}