11use core:: { alloc:: Layout , convert:: TryInto } ;
2+ use rand:: {
3+ distributions:: { Distribution , Uniform } ,
4+ seq:: IteratorRandom ,
5+ } ;
6+ use rand_chacha:: ChaCha20Rng ;
27use usize_conversions:: IntoUsize ;
38use x86_64:: {
49 structures:: paging:: { Page , PageTableIndex , Size4KiB } ,
@@ -7,15 +12,19 @@ use x86_64::{
712use xmas_elf:: program:: ProgramHeader ;
813
914use crate :: {
10- binary:: { MemoryRegion , CONFIG } ,
15+ binary:: { entropy , MemoryRegion , CONFIG } ,
1116 BootInfo ,
1217} ;
1318
1419/// Keeps track of used entries in a level 4 page table.
1520///
1621/// Useful for determining a free virtual memory block, e.g. for mapping additional data.
1722pub struct UsedLevel4Entries {
18- entry_state : [ bool ; 512 ] , // whether an entry is in use by the kernel
23+ /// Whether an entry is in use by the kernel.
24+ entry_state : [ bool ; 512 ] ,
25+ /// A random number generator that should be used to generate random addresses or
26+ /// `None` if aslr is disabled.
27+ rng : Option < ChaCha20Rng > ,
1928}
2029
2130impl UsedLevel4Entries {
@@ -25,6 +34,7 @@ impl UsedLevel4Entries {
2534 pub fn new ( max_phys_addr : PhysAddr , regions_len : usize , framebuffer_size : usize ) -> Self {
2635 let mut used = UsedLevel4Entries {
2736 entry_state : [ false ; 512 ] ,
37+ rng : CONFIG . aslr . then ( entropy:: build_rng) ,
2838 } ;
2939
3040 used. entry_state [ 0 ] = true ; // TODO: Can we do this dynamically?
@@ -99,28 +109,58 @@ impl UsedLevel4Entries {
99109 }
100110 }
101111
102- /// Returns a unused level 4 entry and marks it as used.
112+ /// Returns a unused level 4 entry and marks it as used. If `CONFIG.aslr` is
113+ /// enabled, this will return a random available entry.
103114 ///
104115 /// Since this method marks each returned index as used, it can be used multiple times
105116 /// to determine multiple unused virtual memory regions.
106117 pub fn get_free_entry ( & mut self ) -> PageTableIndex {
107- let ( idx, entry) = self
118+ // Create an iterator over all available p4 indices.
119+ let mut free_entries = self
108120 . entry_state
109- . iter_mut ( )
121+ . iter ( )
122+ . copied ( )
110123 . enumerate ( )
111- . find ( |( _, & mut entry) | entry == false )
112- . expect ( "no usable level 4 entries found" ) ;
124+ . filter ( |( _, used) | !used)
125+ . map ( |( idx, _) | idx) ;
126+
127+ // Choose the free entry index.
128+ let idx_opt = if let Some ( rng) = self . rng . as_mut ( ) {
129+ // Randomly choose an index.
130+ free_entries. choose ( rng)
131+ } else {
132+ // Choose the first index.
133+ free_entries. next ( )
134+ } ;
135+ let idx = idx_opt. expect ( "no usable level 4 entry found" ) ;
136+
137+ // Mark the entry as used.
138+ self . entry_state [ idx] = true ;
113139
114- * entry = true ;
115140 PageTableIndex :: new ( idx. try_into ( ) . unwrap ( ) )
116141 }
117142
118- /// Returns the virtual start address of an unused level 4 entry and marks it as used.
143+ /// Returns a virtual address in an unused level 4 entry and marks it as used.
119144 ///
120- /// This is a convenience method around [`get_free_entry`], so all of its docs applies here
145+ /// This functions call [`get_free_entry`] internally , so all of its docs applies here
121146 /// too.
122- pub fn get_free_address ( & mut self ) -> VirtAddr {
123- Page :: from_page_table_indices_1gib ( self . get_free_entry ( ) , PageTableIndex :: new ( 0 ) )
124- . start_address ( )
147+ pub fn get_free_address ( & mut self , size : u64 , alignment : u64 ) -> VirtAddr {
148+ assert ! ( alignment. is_power_of_two( ) ) ;
149+
150+ let base =
151+ Page :: from_page_table_indices_1gib ( self . get_free_entry ( ) , PageTableIndex :: new ( 0 ) )
152+ . start_address ( ) ;
153+
154+ let offset = if let Some ( rng) = self . rng . as_mut ( ) {
155+ // Choose a random offset.
156+ const LEVEL_4_SIZE : u64 = 4096 * 512 * 512 * 512 ;
157+ let end = LEVEL_4_SIZE - size;
158+ let uniform_range = Uniform :: from ( 0 ..end / alignment) ;
159+ uniform_range. sample ( rng) * alignment
160+ } else {
161+ 0
162+ } ;
163+
164+ base + offset
125165 }
126166}
0 commit comments