x86_64/instructions/
tlb.rs

1//! Functions to flush the translation lookaside buffer (TLB).
2
3use bit_field::BitField;
4
5use crate::{
6    instructions::segmentation::{Segment, CS},
7    structures::paging::{
8        page::{NotGiantPageSize, PageRange},
9        Page, PageSize, Size2MiB, Size4KiB,
10    },
11    PrivilegeLevel, VirtAddr,
12};
13use core::{arch::asm, cmp, convert::TryFrom, fmt};
14
15/// Invalidate the given address in the TLB using the `invlpg` instruction.
16#[inline]
17pub fn flush(addr: VirtAddr) {
18    unsafe {
19        asm!("invlpg [{}]", in(reg) addr.as_u64(), options(nostack, preserves_flags));
20    }
21}
22
23/// Invalidate the TLB completely by reloading the CR3 register.
24#[inline]
25pub fn flush_all() {
26    use crate::registers::control::Cr3;
27    let (frame, flags) = Cr3::read();
28    unsafe { Cr3::write(frame, flags) }
29}
30
31/// The Invalidate PCID Command to execute.
32#[derive(Debug)]
33pub enum InvPcidCommand {
34    /// The logical processor invalidates mappings—except global translations—for the linear address and PCID specified.
35    Address(VirtAddr, Pcid),
36
37    /// The logical processor invalidates all mappings—except global translations—associated with the PCID.
38    Single(Pcid),
39
40    /// The logical processor invalidates all mappings—including global translations—associated with any PCID.
41    All,
42
43    /// The logical processor invalidates all mappings—except global translations—associated with any PCID.
44    AllExceptGlobal,
45}
46
47// TODO: Remove this in the next breaking release.
48#[deprecated = "please use `InvPcidCommand` instead"]
49#[doc(hidden)]
50pub type InvPicdCommand = InvPcidCommand;
51
52/// The INVPCID descriptor comprises 128 bits and consists of a PCID and a linear address.
53/// For INVPCID type 0, the processor uses the full 64 bits of the linear address even outside 64-bit mode; the linear address is not used for other INVPCID types.
54#[repr(C)]
55#[derive(Debug)]
56struct InvpcidDescriptor {
57    pcid: u64,
58    address: u64,
59}
60
61/// Structure of a PCID. A PCID has to be <= 4096 for x86_64.
62#[repr(transparent)]
63#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
64pub struct Pcid(u16);
65
66impl Pcid {
67    /// Create a new PCID. Will result in a failure if the value of
68    /// PCID is out of expected bounds.
69    pub const fn new(pcid: u16) -> Result<Pcid, PcidTooBig> {
70        if pcid >= 4096 {
71            Err(PcidTooBig(pcid))
72        } else {
73            Ok(Pcid(pcid))
74        }
75    }
76
77    /// Get the value of the current PCID.
78    pub const fn value(&self) -> u16 {
79        self.0
80    }
81}
82
83/// A passed `u16` was not a valid PCID.
84///
85/// A PCID has to be <= 4096 for x86_64.
86#[derive(Debug)]
87pub struct PcidTooBig(u16);
88
89impl fmt::Display for PcidTooBig {
90    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
91        write!(f, "PCID should be < 4096, got {}", self.0)
92    }
93}
94
95/// Invalidate the given address in the TLB using the `invpcid` instruction.
96///
97/// ## Safety
98///
99/// This function is unsafe as it requires CPUID.(EAX=07H, ECX=0H):EBX.INVPCID to be 1.
100#[inline]
101pub unsafe fn flush_pcid(command: InvPcidCommand) {
102    let mut desc = InvpcidDescriptor {
103        pcid: 0,
104        address: 0,
105    };
106
107    let kind: u64;
108    match command {
109        InvPcidCommand::Address(addr, pcid) => {
110            kind = 0;
111            desc.pcid = pcid.value().into();
112            desc.address = addr.as_u64()
113        }
114        InvPcidCommand::Single(pcid) => {
115            kind = 1;
116            desc.pcid = pcid.0.into()
117        }
118        InvPcidCommand::All => kind = 2,
119        InvPcidCommand::AllExceptGlobal => kind = 3,
120    }
121
122    unsafe {
123        asm!("invpcid {0}, [{1}]", in(reg) kind, in(reg) &desc, options(nostack, preserves_flags));
124    }
125}
126
127/// Used to broadcast flushes to all logical processors.
128///
129/// ```no_run
130/// use x86_64::VirtAddr;
131/// use x86_64::structures::paging::Page;
132/// use x86_64::instructions::tlb::Invlpgb;
133///
134/// // Check that `invlpgb` and `tlbsync` are supported.
135/// let invlpgb = Invlpgb::new().unwrap();
136///
137/// // Broadcast flushing some pages to all logical processors.
138/// let start: Page = Page::from_start_address(VirtAddr::new(0xf000_0000)).unwrap();
139/// let pages = Page::range(start, start + 3);
140/// invlpgb.build().pages(pages).include_global().flush();
141///
142/// // Wait for all logical processors to respond.
143/// invlpgb.tlbsync();
144/// ```
145#[derive(Debug, Clone, Copy)]
146pub struct Invlpgb {
147    invlpgb_count_max: u16,
148    tlb_flush_nested: bool,
149    nasid: u32,
150}
151
152impl Invlpgb {
153    /// Check that `invlpgb` and `tlbsync` are supported and query limits.
154    ///
155    /// # Panics
156    ///
157    /// Panics if the CPL is not 0.
158    pub fn new() -> Option<Self> {
159        let cs = CS::get_reg();
160        assert_eq!(cs.rpl(), PrivilegeLevel::Ring0);
161
162        // Check if the `INVLPGB` and `TLBSYNC` instruction are supported.
163        let cpuid = unsafe { core::arch::x86_64::__cpuid(0x8000_0008) };
164        if !cpuid.ebx.get_bit(3) {
165            return None;
166        }
167
168        let tlb_flush_nested = cpuid.ebx.get_bit(21);
169        let invlpgb_count_max = cpuid.edx.get_bits(0..=15) as u16;
170
171        // Figure out the number of supported ASIDs.
172        let cpuid = unsafe { core::arch::x86_64::__cpuid(0x8000_000a) };
173        let nasid = cpuid.ebx;
174
175        Some(Self {
176            tlb_flush_nested,
177            invlpgb_count_max,
178            nasid,
179        })
180    }
181
182    /// Returns the maximum count of pages to be flushed supported by the processor.
183    #[inline]
184    pub fn invlpgb_count_max(&self) -> u16 {
185        self.invlpgb_count_max
186    }
187
188    /// Returns whether the processor supports flushing translations used for guest translation.
189    #[inline]
190    pub fn tlb_flush_nested(&self) -> bool {
191        self.tlb_flush_nested
192    }
193
194    /// Returns the number of available address space identifiers.
195    #[inline]
196    pub fn nasid(&self) -> u32 {
197        self.nasid
198    }
199
200    /// Create a `InvlpgbFlushBuilder`.
201    pub fn build(&self) -> InvlpgbFlushBuilder<'_> {
202        InvlpgbFlushBuilder {
203            invlpgb: self,
204            page_range: None,
205            pcid: None,
206            asid: None,
207            include_global: false,
208            final_translation_only: false,
209            include_nested_translations: false,
210        }
211    }
212
213    /// Wait for all previous `invlpgb` instruction executed on the current
214    /// logical processor to be acknowledged by all other logical processors.
215    #[inline]
216    pub fn tlbsync(&self) {
217        unsafe {
218            asm!("tlbsync", options(nomem, preserves_flags));
219        }
220    }
221}
222
223/// A builder struct to construct the parameters for the `invlpgb` instruction.
224#[derive(Debug, Clone)]
225#[must_use]
226pub struct InvlpgbFlushBuilder<'a, S = Size4KiB>
227where
228    S: NotGiantPageSize,
229{
230    invlpgb: &'a Invlpgb,
231    page_range: Option<PageRange<S>>,
232    pcid: Option<Pcid>,
233    asid: Option<u16>,
234    include_global: bool,
235    final_translation_only: bool,
236    include_nested_translations: bool,
237}
238
239impl<'a, S> InvlpgbFlushBuilder<'a, S>
240where
241    S: NotGiantPageSize,
242{
243    /// Flush a range of pages.
244    ///
245    /// If the range doesn't fit within `invlpgb_count_max`, `invlpgb` is
246    /// executed multiple times.
247    pub fn pages<T>(self, page_range: PageRange<T>) -> InvlpgbFlushBuilder<'a, T>
248    where
249        T: NotGiantPageSize,
250    {
251        InvlpgbFlushBuilder {
252            invlpgb: self.invlpgb,
253            page_range: Some(page_range),
254            pcid: self.pcid,
255            asid: self.asid,
256            include_global: self.include_global,
257            final_translation_only: self.final_translation_only,
258            include_nested_translations: self.include_nested_translations,
259        }
260    }
261
262    /// Only flush TLB entries with the given PCID.
263    ///
264    /// # Safety
265    ///
266    /// The caller has to ensure that PCID is enabled in CR4 when the flush is executed.
267    pub unsafe fn pcid(&mut self, pcid: Pcid) -> &mut Self {
268        self.pcid = Some(pcid);
269        self
270    }
271
272    /// Only flush TLB entries with the given ASID.
273    ///
274    /// # Safety
275    ///
276    /// The caller has to ensure that SVM is enabled in EFER when the flush is executed.
277    // FIXME: Make ASID a type and remove error type.
278    pub unsafe fn asid(&mut self, asid: u16) -> Result<&mut Self, AsidOutOfRangeError> {
279        if u32::from(asid) >= self.invlpgb.nasid {
280            return Err(AsidOutOfRangeError {
281                asid,
282                nasid: self.invlpgb.nasid,
283            });
284        }
285
286        self.asid = Some(asid);
287        Ok(self)
288    }
289
290    /// Also flush global pages.
291    pub fn include_global(&mut self) -> &mut Self {
292        self.include_global = true;
293        self
294    }
295
296    /// Only flush the final translation and not the cached upper level TLB entries.
297    pub fn final_translation_only(&mut self) -> &mut Self {
298        self.final_translation_only = true;
299        self
300    }
301
302    /// Also flush nestred translations that could be used for guest translation.
303    pub fn include_nested_translations(mut self) -> Self {
304        assert!(
305            self.invlpgb.tlb_flush_nested,
306            "flushing all nested translations is not supported"
307        );
308
309        self.include_nested_translations = true;
310        self
311    }
312
313    /// Execute the flush.
314    pub fn flush(&self) {
315        if let Some(mut pages) = self.page_range {
316            while !pages.is_empty() {
317                // Calculate out how many pages we still need to flush.
318                let count = Page::<S>::steps_between_impl(&pages.start, &pages.end).0;
319
320                // Make sure that we never jump the gap in the address space when flushing.
321                let second_half_start =
322                    Page::<S>::containing_address(VirtAddr::new(0xffff_8000_0000_0000));
323                let count = if pages.start < second_half_start {
324                    let count_to_second_half =
325                        Page::steps_between_impl(&pages.start, &second_half_start).0;
326                    cmp::min(count, count_to_second_half)
327                } else {
328                    count
329                };
330
331                // We can flush at most u16::MAX pages at once.
332                let count = u16::try_from(count).unwrap_or(u16::MAX);
333
334                // Cap the count by the maximum supported count of the processor.
335                let count = cmp::min(count, self.invlpgb.invlpgb_count_max);
336
337                unsafe {
338                    flush_broadcast(
339                        Some((pages.start, count)),
340                        self.pcid,
341                        self.asid,
342                        self.include_global,
343                        self.final_translation_only,
344                        self.include_nested_translations,
345                    );
346                }
347
348                // Even if the count is zero, one page is still flushed and so
349                // we need to advance by at least one.
350                let inc_count = cmp::max(count, 1);
351                pages.start =
352                    Page::forward_checked_impl(pages.start, usize::from(inc_count)).unwrap();
353            }
354        } else {
355            unsafe {
356                flush_broadcast::<S>(
357                    None,
358                    self.pcid,
359                    self.asid,
360                    self.include_global,
361                    self.final_translation_only,
362                    self.include_nested_translations,
363                );
364            }
365        }
366    }
367}
368
369/// An error returned when trying to use an invalid ASID.
370#[derive(Debug)]
371pub struct AsidOutOfRangeError {
372    /// The requested ASID.
373    pub asid: u16,
374    /// The number of valid ASIDS.
375    pub nasid: u32,
376}
377
378impl fmt::Display for AsidOutOfRangeError {
379    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
380        write!(
381            f,
382            "{} is out of the range of available ASIDS ({})",
383            self.asid, self.nasid
384        )
385    }
386}
387
388/// See `INVLPGB` in AMD64 Architecture Programmer's Manual Volume 3
389#[inline]
390unsafe fn flush_broadcast<S>(
391    va_and_count: Option<(Page<S>, u16)>,
392    pcid: Option<Pcid>,
393    asid: Option<u16>,
394    include_global: bool,
395    final_translation_only: bool,
396    include_nested_translations: bool,
397) where
398    S: NotGiantPageSize,
399{
400    let mut rax = 0;
401    let mut ecx = 0;
402    let mut edx = 0;
403
404    if let Some((va, count)) = va_and_count {
405        rax.set_bit(0, true);
406        rax.set_bits(12.., va.start_address().as_u64().get_bits(12..));
407
408        ecx.set_bits(0..=15, u32::from(count));
409        ecx.set_bit(31, S::SIZE == Size2MiB::SIZE);
410    }
411
412    if let Some(pcid) = pcid {
413        rax.set_bit(1, true);
414        edx.set_bits(16..=27, u32::from(pcid.value()));
415    }
416
417    if let Some(asid) = asid {
418        rax.set_bit(2, true);
419        edx.set_bits(0..=15, u32::from(asid));
420    }
421
422    rax.set_bit(3, include_global);
423    rax.set_bit(4, final_translation_only);
424    rax.set_bit(5, include_nested_translations);
425
426    unsafe {
427        asm!(
428            "invlpgb",
429            in("rax") rax,
430            in("ecx") ecx,
431            in("edx") edx,
432            options(nostack, preserves_flags),
433        );
434    }
435}