strat9_kernel/hardware/storage/
nvme.rs

1// NVMe block device driver
2// Reference: NVM Express Base Specification 2.0
3//
4// Features:
5// - Controller initialization and admin queue
6// - I/O queue pair with IRQ-driven completion
7// - NVMe read (opcode 0x02) and write (opcode 0x01) commands
8// - WaitQueue-based synchronous I/O
9// - Namespace identification
10
11#![allow(dead_code)]
12
13use crate::{
14    arch::x86_64::boot_timestamp,
15    hardware::pci_client::{self as pci, Bar, ProbeCriteria},
16    memory::{allocate_zeroed_frame, paging, phys_to_virt},
17    sync::waitqueue::WaitQueue,
18};
19use alloc::{boxed::Box, format, string::String, sync::Arc, vec::Vec};
20use core::{
21    ptr,
22    sync::atomic::{AtomicBool, AtomicU8, Ordering},
23};
24use spin::Mutex;
25
26const NVME_PAGE_SIZE: usize = 4096;
27const IO_QUEUE_SIZE: usize = 64;
28const MAX_IO_COMMANDS: usize = IO_QUEUE_SIZE;
29const MAX_PRP_ENTRIES: usize = NVME_PAGE_SIZE / 8; // 512 entries per PRP list page
30
31const ADMIN_CQE_ERROR: u16 = (0x1 << 14) | (0x1 << 10);
32
33// =========================================================================
34// TSC-based timeout helpers
35// =========================================================================
36
37/// Read the TSC (Time Stamp Counter).
38#[inline]
39fn rdtsc() -> u64 {
40    unsafe { crate::arch::rdtsc() }
41}
42
43/// Return the calibrated TSC frequency in KHz, or a fallback (3 GHz).
44fn tsc_khz() -> u64 {
45    let khz = boot_timestamp::tsc_khz();
46    if khz > 0 {
47        khz
48    } else {
49        3_000_000
50    }
51}
52
53/// Convert TSC ticks to approximate milliseconds using calibrated frequency.
54fn tsc_to_ms(ticks: u64) -> u64 {
55    ticks / tsc_khz()
56}
57
58/// Get a TSC deadline for `ms` milliseconds from now.
59fn tsc_deadline_ms(ms: u32) -> u64 {
60    rdtsc().wrapping_add((ms as u64) * tsc_khz())
61}
62
63/// Check if a TSC deadline has expired.
64fn tsc_expired(deadline: u64) -> bool {
65    rdtsc() >= deadline
66}
67
68/// Read a 64-bit register from the NVMe controller.
69unsafe fn regs_read64(base: usize, offset: u64) -> u64 {
70    let low = core::ptr::read_volatile((base + offset as usize) as *const u32) as u64;
71    let high = core::ptr::read_volatile((base + offset as usize + 4) as *const u32) as u64;
72    low | (high << 32)
73}
74
75/// Build PRP (Physical Region Page) list for a physically contiguous buffer.
76///
77/// Returns (prp1, prp2):
78/// - If the buffer fits in a single page: prp1 = buf_phys, prp2 = 0
79/// - If the buffer spans multiple pages: prp1 = buf_phys, prp2 = PRP list phys addr
80///
81/// The caller must ensure `buf_phys` is page-aligned and the buffer is physically
82/// contiguous for multi-page transfers.
83unsafe fn build_prp_list(buf_phys: u64, byte_count: usize) -> (u64, u64) {
84    let first_page_remaining = NVME_PAGE_SIZE - (buf_phys as usize % NVME_PAGE_SIZE);
85    if byte_count <= first_page_remaining {
86        return (buf_phys, 0);
87    }
88
89    // Allocate a PRP list page
90    let prp_frame = allocate_zeroed_frame().expect("NVMe: failed to allocate PRP list");
91    let prp_phys = prp_frame.start_address.as_u64();
92    paging::ensure_identity_map_range(prp_phys, NVME_PAGE_SIZE as u64);
93    let prp_virt = phys_to_virt(prp_phys) as *mut u64;
94
95    // First page covers buf_phys..end_of_first_page
96    let mut remaining = byte_count - first_page_remaining;
97    let mut next_phys = (buf_phys & !0xFFF) + NVME_PAGE_SIZE as u64;
98    let mut idx = 0;
99
100    while remaining > 0 && idx < MAX_PRP_ENTRIES {
101        core::ptr::write_volatile(prp_virt.add(idx), next_phys);
102        idx += 1;
103        next_phys += NVME_PAGE_SIZE as u64;
104        remaining = remaining.saturating_sub(NVME_PAGE_SIZE);
105    }
106
107    (buf_phys, prp_phys)
108}
109
110#[repr(transparent)]
111struct VolatileCell<T> {
112    value: T,
113}
114
115impl<T> VolatileCell<T> {
116    fn read(&self) -> T
117    where
118        T: Copy,
119    {
120        unsafe { ptr::read_volatile(&self.value) }
121    }
122    fn write(&self, val: T) {
123        unsafe { ptr::write_volatile(core::ptr::addr_of!(self.value) as *mut T, val) }
124    }
125}
126
127unsafe impl<T: Send> Send for VolatileCell<T> {}
128unsafe impl<T: Sync> Sync for VolatileCell<T> {}
129
130#[repr(C)]
131struct Capability {
132    value: VolatileCell<u64>,
133}
134
135impl Capability {
136    fn max_queue_entries(&self) -> u16 {
137        (self.value.read() & 0xFFFF) as u16
138    }
139    fn doorbell_stride(&self) -> u64 {
140        (self.value.read() >> 32) & 0xF
141    }
142}
143
144#[repr(transparent)]
145struct Version {
146    value: VolatileCell<u32>,
147}
148
149#[repr(C)]
150struct ControllerConfig {
151    value: VolatileCell<u32>,
152}
153
154impl ControllerConfig {
155    fn clear_io_fields(&self) {
156        let mut val = self.value.read();
157        val &= !(((0xF) << 16) | ((0xF) << 20) | ((0x7) << 4));
158        self.value.write(val);
159    }
160    fn set_iosqes(&self, size: u32) {
161        let mut val = self.value.read();
162        val |= (size & 0xF) << 16;
163        self.value.write(val);
164    }
165    fn set_iocqes(&self, size: u32) {
166        let mut val = self.value.read();
167        val |= (size & 0xF) << 20;
168        self.value.write(val);
169    }
170    fn set_css(&self, css: u32) {
171        let mut val = self.value.read();
172        val |= (css & 0x7) << 4;
173        self.value.write(val);
174    }
175    fn set_enable(&self, enable: bool) {
176        let mut val = self.value.read();
177        if enable {
178            val |= 1;
179        } else {
180            val &= !1;
181        }
182        self.value.write(val);
183    }
184    fn is_enabled(&self) -> bool {
185        (self.value.read() & 1) != 0
186    }
187}
188
189#[repr(transparent)]
190struct ControllerStatus {
191    value: VolatileCell<u32>,
192}
193
194impl ControllerStatus {
195    fn is_ready(&self) -> bool {
196        (self.value.read() & 1) != 0
197    }
198    fn is_fatal(&self) -> bool {
199        (self.value.read() >> 1) & 1 != 0
200    }
201}
202
203#[repr(C)]
204struct Registers {
205    capability: Capability,
206    version: Version,
207    _intms: VolatileCell<u32>,
208    _intmc: VolatileCell<u32>,
209    cc: ControllerConfig,
210    _reserved1: VolatileCell<u32>,
211    csts: ControllerStatus,
212    _reserved2: VolatileCell<u32>,
213    aqa: VolatileCell<u32>,
214    asq_low: VolatileCell<u32>,
215    asq_high: VolatileCell<u32>,
216    acq_low: VolatileCell<u32>,
217    acq_high: VolatileCell<u32>,
218}
219
220#[derive(Debug, Clone, Copy)]
221enum NvmeError {
222    ControllerFatal,
223    Timeout,
224    InvalidNamespace,
225    IoError,
226}
227
228#[derive(Debug, Clone)]
229pub struct NvmeNamespace {
230    pub nsid: u32,
231    pub size: u64,
232    pub block_size: u32,
233}
234
235/// LBA Format descriptor (4 bytes each, starting at byte 128 of Identify Namespace).
236#[repr(C)]
237#[derive(Copy, Clone)]
238struct LbaFormat {
239    _metadata_size: u16,
240    lbads: u8,
241    _relative_perf: u8,
242}
243
244/// Identify Namespace data (first 384 bytes needed).
245/// Layout per NVMe spec: bytes 0-73 are explicit fields, 74-127 reserved, 128-383 LBAF[0..64].
246#[repr(C)]
247#[derive(Copy, Clone)]
248struct IdentifyNamespaceData {
249    nsze: u64,
250    _ncap: u64,
251    _nuse: u64,
252    _nsfeat: u8,
253    _nlbaf: u8,
254    flbas: u8,
255    _mc: u8,
256    _dpc: u8,
257    _dps: u8,
258    _nmic: u8,
259    _rescap: u8,
260    _fpi: u8,
261    _dlfeat: u8,
262    _nawun: u16,
263    _nawupf: u16,
264    _nacwu: u16,
265    _nabsn: u16,
266    _nabo: u16,
267    _nabspf: u16,
268    _noiob: u16,
269    _nvmcap: [u64; 2],
270    _npwg: u16,
271    _npwa: u16,
272    _npdg: u16,
273    _npda: u16,
274    _nows: u16,
275    _reserved: [u8; 54], // bytes 74-127
276    /// LBA Format Support : 64 entries × 4 bytes at offset 128
277    lbaf: [LbaFormat; 64],
278}
279
280/// Identify Controller data (partial : fields we need).
281#[repr(C)]
282#[derive(Copy, Clone)]
283struct IdentifyControllerData {
284    _vid: u16,
285    _ssvid: u16,
286    _sn: [u8; 20],
287    _mn: [u8; 40],
288    _fr: [u8; 8],
289    _rab: u8,
290    _ieee: [u8; 3],
291    _cmic: u8,
292    mdts: u8,
293    _cntlid: u16,
294    _ver: u32,
295    _rtd3r: u32,
296    _rtd3e: u32,
297    _oaes: u32,
298    _ctratt: u32,
299    _reserved0: [u8; 100],
300    _oacs: u16,
301    _acl: u8,
302    _aerl: u8,
303    _frmw: u8,
304    _lpa: u8,
305    _elpe: u8,
306    _npss: u8,
307    _avscc: u8,
308    _apsta: u8,
309    _wctemp: u16,
310    _cctemp: u16,
311    _reserved1: [u8; 242],
312    sqes: u8,
313    cqes: u8,
314    _maxcmd: u16,
315    nn: u32,
316    _oncs: u16,
317    _fuses: u16,
318    _fna: u8,
319    _vwc: u8,
320    _awun: u16,
321    _awupf: u16,
322    _icsvscc: u8,
323    _nwpc: u8,
324    _acwu: u16,
325    _cdfs: u16,
326    _sgls: u32,
327    _reserved2: [u8; 228],
328    _subnqn: [u8; 256],
329    _reserved3: [u8; 1024],
330    _psd: [[u8; 32]; 32],
331    _vs: [u8; 1024],
332}
333
334struct IoQueuePair {
335    submission: IoQueue<Submission>,
336    completion: IoQueue<Completion>,
337    command_id: u16,
338    size: usize,
339}
340
341struct IoQueue<T: QueueType> {
342    doorbell: *const VolatileCell<u32>,
343    entries: *mut T::EntryType,
344    size: usize,
345    index: usize,
346    phase: bool,
347    phys_addr: u64,
348}
349
350unsafe impl<T: QueueType> Send for IoQueue<T> {}
351unsafe impl<T: QueueType> Sync for IoQueue<T> {}
352
353impl<T: QueueType> IoQueue<T> {
354    fn new(registers_base: usize, size: usize, queue_id: u16, dstrd: usize) -> Self {
355        let doorbell_offset =
356            0x1000 + ((((queue_id as usize) * 2) + T::DOORBELL_OFFSET) * (4 << dstrd));
357        let doorbell =
358            unsafe { &*((registers_base + doorbell_offset) as *const VolatileCell<u32>) };
359
360        let frame = allocate_zeroed_frame().expect("NVMe: failed to allocate I/O queue frame");
361        let phys_addr = frame.start_address.as_u64();
362        paging::ensure_identity_map_range(phys_addr, NVME_PAGE_SIZE as u64);
363        let virt_addr = phys_to_virt(phys_addr);
364
365        unsafe {
366            ptr::write_bytes(
367                virt_addr as *mut u8,
368                0,
369                size * core::mem::size_of::<T::EntryType>(),
370            );
371        }
372
373        Self {
374            doorbell,
375            entries: virt_addr as *mut T::EntryType,
376            size,
377            index: 0,
378            phase: true,
379            phys_addr,
380        }
381    }
382}
383
384#[derive(Debug, Clone, Copy)]
385pub struct NvmeCompletionResult {
386    pub command_id: u16,
387    pub status: u16,
388}
389
390pub struct NvmeController {
391    registers: usize,
392    admin_queue: Mutex<QueuePair>,
393    io_queue: Mutex<IoQueuePair>,
394    namespaces: Vec<NvmeNamespace>,
395    pub name: String,
396    irq_line: u8,
397    io_done: Box<[AtomicBool]>,
398    io_wq: WaitQueue,
399}
400
401unsafe impl Send for NvmeController {}
402unsafe impl Sync for NvmeController {}
403
404impl NvmeController {
405    unsafe fn new(registers: usize, name: String) -> Result<Self, NvmeError> {
406        let regs = &*(registers as *const Registers);
407        let dstrd = regs.capability.doorbell_stride() as usize;
408        let max_entries = regs.capability.max_queue_entries();
409        let queue_size = core::cmp::min(max_entries as usize, 1024);
410
411        let admin_queue = QueuePair::new(registers, queue_size, dstrd);
412
413        let io_sub = IoQueue::new(registers, IO_QUEUE_SIZE, 1, dstrd);
414        let io_comp = IoQueue::new(registers, IO_QUEUE_SIZE, 1, dstrd);
415        let io_queue = IoQueuePair {
416            submission: io_sub,
417            completion: io_comp,
418            command_id: 0,
419            size: IO_QUEUE_SIZE,
420        };
421
422        let io_done: Box<[AtomicBool]> = (0..MAX_IO_COMMANDS)
423            .map(|_| AtomicBool::new(false))
424            .collect();
425
426        let mut controller = Self {
427            registers,
428            admin_queue: Mutex::new(admin_queue),
429            io_queue: Mutex::new(io_queue),
430            namespaces: Vec::new(),
431            name,
432            irq_line: 0,
433            io_done,
434            io_wq: WaitQueue::new(),
435        };
436
437        controller.init_admin_queue()?;
438        controller.create_io_queues()?;
439        controller.identify_namespaces()?;
440        Ok(controller)
441    }
442
443    fn submit_admin_command(&self, command: Command) -> Result<CompletionEntry, NvmeError> {
444        let mut admin = self.admin_queue.lock();
445        admin.submit_command(command).ok_or(NvmeError::IoError)
446    }
447
448    fn init_admin_queue(&mut self) -> Result<(), NvmeError> {
449        let regs = unsafe { &*(self.registers as *const Registers) };
450        let (admin_sq_phys, admin_cq_phys, queue_size) = {
451            let q = self.admin_queue.lock();
452            (q.submission_phys(), q.completion_phys(), q.size)
453        };
454
455        if queue_size == 0 {
456            return Err(NvmeError::IoError);
457        }
458        let qsz = ((queue_size as u32).saturating_sub(1)) & 0x0FFF;
459
460        log::info!("NVMe init: queue_size={} qsz={}", queue_size, qsz);
461        log::info!(
462            "NVMe init: ASQ phys={:#x} ACQ phys={:#x}",
463            admin_sq_phys,
464            admin_cq_phys
465        );
466
467        // Step 1: Disable controller if it's already enabled
468        let cc_val = regs.cc.value.read();
469        let csts_val = regs.csts.value.read();
470        log::info!("NVMe init: CC={:#010x} CSTS={:#010x}", cc_val, csts_val);
471
472        if cc_val & 1 != 0 {
473            log::info!("NVMe init: controller enabled, disabling...");
474
475            // Some AMD NVMe controllers (e.g. on Lenovo X13) hang if we send
476            // SHN (shutdown notification) before clearing CC.EN.  The safe
477            // approach for re-initialisation is to skip SHN and just write
478            // CC.EN=0 directly.  This matches Linux's nvme_disable_ctrl().
479            regs.cc.set_enable(false);
480            log::info!("NVMe init: wrote CC.EN=0 (no SHN), waiting for CSTS.RDY...");
481
482            // Wait for CSTS.RDY (bit 0) to clear : up to 5.5 s per NVMe spec.
483            // Use a generous timeout; some AMD controllers are slow to respond.
484            let deadline = tsc_deadline_ms(5500);
485            let mut log_count = 0u32;
486            loop {
487                let csts = regs.csts.value.read();
488                if csts & 1 == 0 {
489                    log::info!("NVMe init: controller disabled, CSTS={:#x}", csts);
490                    break;
491                }
492                if tsc_expired(deadline) {
493                    log::warn!("NVMe init: RDY timeout, CSTS={:#x} : forcing CC=0", csts);
494
495                    // Last resort: write CC=0x00000000
496                    regs.cc.value.write(0x0000_0000);
497                    log::info!("NVMe init: wrote CC=0x00000000 (full reset)");
498                    let force_deadline = tsc_deadline_ms(3000);
499                    while !tsc_expired(force_deadline) {
500                        let c = regs.csts.value.read();
501                        if c & 1 == 0 {
502                            log::info!("NVMe init: forced disable OK, CSTS={:#x}", c);
503                            break;
504                        }
505                        core::hint::spin_loop();
506                    }
507                    break;
508                }
509                core::hint::spin_loop();
510                log_count += 1;
511                if log_count % 2_000_000 == 0 {
512                    log::info!(
513                        "NVMe init: still waiting... CSTS={:#x}",
514                        regs.csts.value.read()
515                    );
516                }
517            }
518        } else {
519            log::info!("NVMe init: controller already disabled");
520        }
521
522        // Small delay after disable ; some AMD controllers need this before
523        // admin queue registers become writable.
524        let settle = tsc_deadline_ms(2);
525        while !tsc_expired(settle) {
526            core::hint::spin_loop();
527        }
528
529        // Step 2: Verify CSTS
530        let csts = regs.csts.value.read();
531        log::info!("NVMe init: CSTS={:#x}", csts);
532        if csts & 2 != 0 {
533            log::error!("NVMe init: CSTS.CFS (fatal) set!");
534            return Err(NvmeError::ControllerFatal);
535        }
536
537        // Step 3: Write admin queue registers (32-bit writes, like Redox/MaestroOS)
538        log::info!("NVMe init: writing AQA={:#x}...", qsz | (qsz << 16));
539        regs.aqa.write(qsz | (qsz << 16));
540        log::info!("NVMe init: writing ASQ={:#x}...", admin_sq_phys);
541        regs.asq_low.write(admin_sq_phys as u32);
542        regs.asq_high.write((admin_sq_phys >> 32) as u32);
543        log::info!("NVMe init: writing ACQ={:#x}...", admin_cq_phys);
544        regs.acq_low.write(admin_cq_phys as u32);
545        regs.acq_high.write((admin_cq_phys >> 32) as u32);
546
547        // Step 4: Configure CC
548        log::info!("NVMe init: configuring CC...");
549        regs.cc.clear_io_fields();
550        regs.cc.set_css(0); // NVM Command Set
551        regs.cc.set_iosqes(6); // 2^6 = 64 byte submission queue entries
552        regs.cc.set_iocqes(6); // 2^6 = 64 byte completion queue entries
553        log::info!("NVMe init: CC={:#010x} (configured)", regs.cc.value.read());
554
555        // Step 5: Enable controller
556        log::info!("NVMe init: enabling controller...");
557        regs.cc.set_enable(true);
558        log::info!("NVMe init: CC={:#010x} (EN=1)", regs.cc.value.read());
559
560        // Step 6: Wait for CSTS.RDY (up to 5.5s per NVMe spec)
561        let deadline = tsc_deadline_ms(5500);
562        let mut log_interval = 0u32;
563        loop {
564            let csts = regs.csts.value.read();
565            if csts & 1 != 0 {
566                log::info!("NVMe init: controller ready, CSTS={:#x}", csts);
567                break;
568            }
569            if tsc_expired(deadline) {
570                log::error!("NVMe init: enable timeout! CSTS={:#x}", csts);
571                return Err(NvmeError::Timeout);
572            }
573            core::hint::spin_loop();
574            log_interval += 1;
575            if log_interval % 500_000 == 0 {
576                log::info!("NVMe init: waiting for ready... CSTS={:#x}", csts);
577            }
578        }
579
580        let csts = regs.csts.value.read();
581        log::info!("NVMe init: final CSTS={:#x}", csts);
582
583        if csts & 2 != 0 {
584            log::error!("NVMe init: controller fatal error!");
585            return Err(NvmeError::ControllerFatal);
586        }
587
588        log::info!(
589            "NVMe: Controller v{}.{}.{} ready",
590            regs.version.value.read() >> 16,
591            (regs.version.value.read() >> 8) & 0xFF,
592            regs.version.value.read() & 0xFF
593        );
594        Ok(())
595    }
596
597    fn create_io_queues(&mut self) -> Result<(), NvmeError> {
598        let (io_sq_phys, io_cq_phys, queue_size) = {
599            let q = self.io_queue.lock();
600            (q.submission.phys_addr, q.completion.phys_addr, q.size)
601        };
602
603        let qsz = ((queue_size as u32).saturating_sub(1)) & 0xFFF;
604
605        log::info!(
606            "NVMe I/O queues: size={} qsz={} SQ={:#x} CQ={:#x}",
607            queue_size,
608            qsz,
609            io_sq_phys,
610            io_cq_phys
611        );
612
613        log::info!("NVMe I/O: sending Set Features (IRQ coalescing)...");
614        let set_feature_cmd = Command {
615            opcode: 0x09,
616            cdw10: 0x07,
617            cdw11: 0x0100_0000,
618            ..Default::default()
619        };
620        self.submit_admin_command(set_feature_cmd).ok();
621
622        log::info!("NVMe I/O: creating I/O CQ...");
623        // NVMe spec CDW10: bits 15:0=QID, bits 31:16=QSIZE(0-based)
624        // NVMe spec CDW11: bit 0=PCIE, bit 1=IEN, bits 31:16=IV
625        let cq_cmd = Command {
626            opcode: 0x05,
627            cdw10: qsz << 16,   // QID=0 in bits 15:0, QSIZE in bits 31:16
628            cdw11: 0x0000_0003, // PCIE=1, IEN=1
629            prp1: io_cq_phys,
630            ..Default::default()
631        };
632        match self.submit_admin_command(cq_cmd) {
633            Ok(c) => {
634                if c.status_code() != 0 {
635                    log::warn!("NVMe: Create I/O CQ failed: status={}", c.status_code());
636                } else {
637                    log::info!("NVMe I/O: CQ created OK");
638                }
639            }
640            Err(e) => {
641                log::warn!("NVMe: Create I/O CQ error: {:?}", e);
642                return Err(e);
643            }
644        }
645
646        log::info!("NVMe I/O: creating I/O SQ...");
647        // NVMe spec CDW10: bits 15:0=QID, bits 31:16=QSIZE(0-based)
648        // NVMe spec CDW11: bit 0=PCIE, bits 31:16=CQID
649        let sq_cmd = Command {
650            opcode: 0x01,
651            cdw10: (qsz << 16) | 1, // QID=1 in bits 15:0, QSIZE in bits 31:16
652            cdw11: (1 << 16) | 0x0000_0001, // CQID=1 in bits 31:16, PCIE=1
653            prp1: io_sq_phys,
654            ..Default::default()
655        };
656        match self.submit_admin_command(sq_cmd) {
657            Ok(c) => {
658                if c.status_code() != 0 {
659                    log::warn!("NVMe: Create I/O SQ failed: status={}", c.status_code());
660                } else {
661                    log::info!("NVMe I/O: SQ created OK");
662                }
663            }
664            Err(e) => {
665                log::warn!("NVMe: Create I/O SQ error: {:?}", e);
666                return Err(e);
667            }
668        }
669
670        log::info!("NVMe: I/O queues created (size={})", queue_size);
671        Ok(())
672    }
673
674    fn identify(&self, cns: u8, nsid: u32) -> Result<Vec<u8>, NvmeError> {
675        let frame = allocate_zeroed_frame().ok_or(NvmeError::IoError)?;
676        let phys = frame.start_address.as_u64();
677        paging::ensure_identity_map_range(phys, NVME_PAGE_SIZE as u64);
678        let virt = phys_to_virt(phys) as *mut u8;
679        unsafe {
680            ptr::write_bytes(virt, 0, NVME_PAGE_SIZE);
681        }
682
683        let cmd = Command {
684            opcode: 0x06,
685            nsid,
686            prp1: phys,
687            cdw10: cns as u32,
688            ..Default::default()
689        };
690
691        let completion = self.submit_admin_command(cmd)?;
692        if completion.status_code() != 0 {
693            return Err(NvmeError::IoError);
694        }
695
696        // Copy data to owned Vec before the DMA frame is freed.
697        let mut data = Vec::with_capacity(NVME_PAGE_SIZE);
698        unsafe {
699            for i in 0..NVME_PAGE_SIZE {
700                data.push(ptr::read_volatile(virt.add(i)));
701            }
702        }
703        // frame dropped here => physical memory freed
704        Ok(data)
705    }
706
707    fn identify_namespaces(&mut self) -> Result<(), NvmeError> {
708        // Step 1: Identify Controller : get NN, MDTS, SQES, CQES, and info
709        let ctrl_data = self.identify(0x01, 0)?;
710        let ctrl = unsafe { core::ptr::read(ctrl_data.as_ptr() as *const IdentifyControllerData) };
711        let nn = ctrl.nn;
712        let mdts = ctrl.mdts;
713
714        // Validate SQES/CQES (like MaestroOS)
715        let min_sqes = ctrl.sqes & 0xF;
716        let max_sqes = (ctrl.sqes >> 4) & 0xF;
717        let min_cqes = ctrl.cqes & 0xF;
718        let max_cqes = (ctrl.cqes >> 4) & 0xF;
719        let our_sqes = 6u8; // 2^6 = 64 bytes
720        let our_cqes = 4u8; // 2^4 = 16 bytes
721
722        if our_sqes < min_sqes || our_sqes > max_sqes {
723            log::warn!(
724                "NVMe: SQES {} not in range [{}..{}] : controller may reject commands",
725                our_sqes,
726                min_sqes,
727                max_sqes
728            );
729        }
730        if our_cqes < min_cqes || our_cqes > max_cqes {
731            log::warn!(
732                "NVMe: CQES {} not in range [{}..{}] : controller may reject completions",
733                our_cqes,
734                min_cqes,
735                max_cqes
736            );
737        }
738
739        // Validate page size against CAP (like MaestroOS)
740        let cap = unsafe { regs_read64(self.registers, 0x00) };
741        let mpsmin = ((cap >> 48) & 0xF) as u32 + 12; // log2 of min page size
742        let mpsmax = ((cap >> 52) & 0xF) as u32 + 12; // log2 of max page size
743        let our_mps = 12u32; // 4096 bytes
744        if our_mps < mpsmin || our_mps > mpsmax {
745            log::warn!(
746                "NVMe: Page size {} not in range [{}..{}] bytes",
747                1 << our_mps,
748                1 << mpsmin,
749                1 << mpsmax
750            );
751        }
752
753        log::info!(
754            "NVMe: Controller NN={} MDTS={} SQES={:#x} CQES={:#x} MPS=[{}..{}]",
755            nn,
756            mdts,
757            ctrl.sqes,
758            ctrl.cqes,
759            1 << mpsmin,
760            1 << mpsmax
761        );
762
763        if nn == 0 {
764            return Err(NvmeError::InvalidNamespace);
765        }
766
767        // Step 2: Identify Active Namespace ID List (CNS=2)
768        let ns_list = self.identify(0x02, 0)?;
769        let ns_list_words = unsafe {
770            core::slice::from_raw_parts(ns_list.as_ptr() as *const u32, ns_list.len() / 4)
771        };
772        let mut active_nsids: Vec<u32> = Vec::new();
773        for &word in ns_list_words.iter() {
774            if word == 0 {
775                break;
776            }
777            active_nsids.push(word);
778        }
779        log::info!("NVMe: {} active namespace(s)", active_nsids.len());
780
781        // Step 3: Identify each active namespace
782        for nsid in &active_nsids {
783            if let Ok(ns_data) = self.identify(0x00, *nsid) {
784                let ns =
785                    unsafe { core::ptr::read(ns_data.as_ptr() as *const IdentifyNamespaceData) };
786
787                let flbas = ns.flbas as usize;
788                let lbaf_idx = flbas & 0xF;
789                let lbads = ns.lbaf[lbaf_idx].lbads;
790                let block_size = 1u32 << lbads;
791
792                self.namespaces.push(NvmeNamespace {
793                    nsid: *nsid,
794                    size: ns.nsze,
795                    block_size,
796                });
797                log::info!(
798                    "NVMe: NSID {} - {} blocks @ {} bytes (LBADS={}, FLBAS={:#x})",
799                    nsid,
800                    ns.nsze,
801                    block_size,
802                    lbads,
803                    flbas
804                );
805            }
806        }
807        Ok(())
808    }
809
810    fn submit_io_command(&self, command: &mut Command) -> Result<u16, NvmeError> {
811        let mut io = self.io_queue.lock();
812        let cmd_id = io.command_id;
813        command.command_id = cmd_id;
814        io.command_id = io.command_id.wrapping_add(1);
815
816        let slot = cmd_id as usize % io.size;
817
818        let idx = cmd_id as usize % MAX_IO_COMMANDS;
819        self.io_done[idx].store(false, Ordering::SeqCst);
820
821        unsafe {
822            ptr::write(io.submission.entries.add(slot), *command);
823            core::sync::atomic::fence(core::sync::atomic::Ordering::SeqCst);
824            (*io.submission.doorbell).write(((slot + 1) % io.size) as u32);
825        }
826
827        Ok(cmd_id)
828    }
829
830    pub fn read_blocks(
831        &self,
832        nsid: u32,
833        lba: u64,
834        block_count: u32,
835        buf_phys: u64,
836    ) -> Result<(), NvmeError> {
837        let byte_count = block_count as usize * 512;
838        let (prp1, prp2) = unsafe { build_prp_list(buf_phys, byte_count) };
839
840        let cmd_id = {
841            let mut cmd = Command {
842                opcode: 0x02,
843                nsid,
844                prp1,
845                prp2,
846                cdw10: (lba & 0xFFFF_FFFF) as u32,
847                cdw11: ((lba >> 32) & 0xFFFF_FFFF) as u32,
848                cdw12: (block_count - 1),
849                ..Default::default()
850            };
851            self.submit_io_command(&mut cmd)?
852        };
853
854        let idx = cmd_id as usize % MAX_IO_COMMANDS;
855
856        self.io_wq.wait_until(|| {
857            if self.io_done[idx].load(Ordering::Acquire) {
858                Some(())
859            } else {
860                None
861            }
862        });
863
864        Ok(())
865    }
866
867    pub fn write_blocks(
868        &self,
869        nsid: u32,
870        lba: u64,
871        block_count: u32,
872        buf_phys: u64,
873    ) -> Result<(), NvmeError> {
874        let byte_count = block_count as usize * 512;
875        let (prp1, prp2) = unsafe { build_prp_list(buf_phys, byte_count) };
876
877        let cmd_id = {
878            let mut cmd = Command {
879                opcode: 0x01,
880                nsid,
881                prp1,
882                prp2,
883                cdw10: (lba & 0xFFFF_FFFF) as u32,
884                cdw11: ((lba >> 32) & 0xFFFF_FFFF) as u32,
885                cdw12: (block_count - 1),
886                ..Default::default()
887            };
888            self.submit_io_command(&mut cmd)?
889        };
890
891        let idx = cmd_id as usize % MAX_IO_COMMANDS;
892
893        self.io_wq.wait_until(|| {
894            if self.io_done[idx].load(Ordering::Acquire) {
895                Some(())
896            } else {
897                None
898            }
899        });
900
901        Ok(())
902    }
903
904    pub fn handle_interrupt(&self) {
905        let mut io = self.io_queue.lock();
906
907        loop {
908            let entry = unsafe { &*io.completion.entries.add(io.completion.index) };
909            let status = entry.status;
910            if ((status & 0x1) != 0) == io.completion.phase {
911                let cmd_id = entry.command_id;
912                let sc = (entry.status >> 1) & 0xFF;
913                let dnr = (entry.status >> 14) & 1;
914
915                io.completion.index = (io.completion.index + 1) % io.completion.size;
916                if io.completion.index == 0 {
917                    io.completion.phase = !io.completion.phase;
918                }
919                unsafe {
920                    (*io.completion.doorbell).write(io.completion.index as u32);
921                }
922
923                let idx = cmd_id as usize % MAX_IO_COMMANDS;
924                if sc != 0 && dnr == 0 {
925                    log::warn!("NVMe: I/O error cmd_id={} sc={}", cmd_id, sc);
926                }
927                self.io_done[idx].store(true, Ordering::Release);
928                self.io_wq.wake_all();
929            } else {
930                break;
931            }
932        }
933    }
934
935    pub fn namespace_count(&self) -> usize {
936        self.namespaces.len()
937    }
938
939    pub fn get_namespace(&self, index: usize) -> Option<&NvmeNamespace> {
940        self.namespaces.get(index)
941    }
942
943    pub fn set_irq_line(&mut self, irq: u8) {
944        self.irq_line = irq;
945    }
946
947    pub fn irq_line(&self) -> u8 {
948        self.irq_line
949    }
950}
951
952#[repr(C)]
953#[derive(Default, Copy, Clone)]
954struct Command {
955    opcode: u8,
956    flags: u8,
957    command_id: u16,
958    nsid: u32,
959    cdw2: u32,
960    cdw3: u32,
961    prp1: u64,
962    prp2: u64,
963    cdw10: u32,
964    cdw11: u32,
965    cdw12: u32,
966    cdw13: u32,
967    cdw14: u32,
968    cdw15: u32,
969}
970
971#[repr(C)]
972#[derive(Copy, Clone)]
973struct CompletionEntry {
974    dw0: u32,
975    dw1: u32,
976    sq_head: u16,
977    sq_id: u16,
978    command_id: u16,
979    status: u16,
980}
981
982impl CompletionEntry {
983    fn status_code(&self) -> u8 {
984        ((self.status >> 1) & 0xFF) as u8
985    }
986}
987
988struct QueuePair {
989    #[allow(dead_code)]
990    id: u16,
991    size: usize,
992    command_id: u16,
993    submission: Queue<Submission>,
994    completion: Queue<Completion>,
995}
996
997struct Submission;
998struct Completion;
999
1000trait QueueType {
1001    type EntryType;
1002    const DOORBELL_OFFSET: usize;
1003}
1004
1005impl QueueType for Submission {
1006    type EntryType = Command;
1007    const DOORBELL_OFFSET: usize = 0;
1008}
1009
1010impl QueueType for Completion {
1011    type EntryType = CompletionEntry;
1012    const DOORBELL_OFFSET: usize = 1;
1013}
1014
1015struct Queue<T: QueueType> {
1016    doorbell: *const VolatileCell<u32>,
1017    entries: *mut T::EntryType,
1018    size: usize,
1019    index: usize,
1020    phase: bool,
1021    phys_addr: u64,
1022}
1023
1024impl<T: QueueType> Queue<T> {
1025    fn new(registers_base: usize, size: usize, queue_id: u16, dstrd: usize) -> Self {
1026        let doorbell_offset =
1027            0x1000 + ((((queue_id as usize) * 2) + T::DOORBELL_OFFSET) * (4 << dstrd));
1028        let doorbell =
1029            unsafe { &*((registers_base + doorbell_offset) as *const VolatileCell<u32>) };
1030
1031        let frame = allocate_zeroed_frame().expect("NVMe: failed to allocate queue frame");
1032        let phys_addr = frame.start_address.as_u64();
1033        paging::ensure_identity_map_range(phys_addr, NVME_PAGE_SIZE as u64);
1034        let virt_addr = phys_to_virt(phys_addr);
1035
1036        unsafe {
1037            ptr::write_bytes(
1038                virt_addr as *mut u8,
1039                0,
1040                size * core::mem::size_of::<T::EntryType>(),
1041            );
1042        }
1043
1044        Self {
1045            doorbell,
1046            entries: virt_addr as *mut T::EntryType,
1047            size,
1048            index: 0,
1049            phase: true,
1050            phys_addr,
1051        }
1052    }
1053
1054    fn phys_addr(&self) -> u64 {
1055        self.phys_addr
1056    }
1057}
1058
1059impl Queue<Completion> {
1060    fn poll_completion(&mut self) -> Option<CompletionEntry> {
1061        unsafe {
1062            let entry = &*self.entries.add(self.index);
1063            let status = entry.status;
1064            if ((status & 0x1) != 0) == self.phase {
1065                let completion = ptr::read(entry);
1066                self.index = (self.index + 1) % self.size;
1067                if self.index == 0 {
1068                    self.phase = !self.phase;
1069                }
1070                (*self.doorbell).write(self.index as u32);
1071                Some(completion)
1072            } else {
1073                None
1074            }
1075        }
1076    }
1077}
1078
1079impl Queue<Submission> {
1080    fn submit_command(&mut self, command: Command, idx: usize) {
1081        unsafe {
1082            ptr::write(self.entries.add(idx), command);
1083            (*self.doorbell).write(((idx + 1) % self.size) as u32);
1084        }
1085        core::sync::atomic::fence(core::sync::atomic::Ordering::SeqCst);
1086    }
1087}
1088
1089impl QueuePair {
1090    fn new(registers_base: usize, size: usize, dstrd: usize) -> Self {
1091        static NEXT_ID: AtomicU8 = AtomicU8::new(0);
1092        let id = NEXT_ID.fetch_add(1, Ordering::SeqCst) as u16;
1093        Self {
1094            id,
1095            size,
1096            command_id: 0,
1097            submission: Queue::new(registers_base, size, id, dstrd),
1098            completion: Queue::new(registers_base, size, id, dstrd),
1099        }
1100    }
1101
1102    fn submission_phys(&self) -> u64 {
1103        self.submission.phys_addr()
1104    }
1105    fn completion_phys(&self) -> u64 {
1106        self.completion.phys_addr()
1107    }
1108
1109    fn submit_command(&mut self, command: Command) -> Option<CompletionEntry> {
1110        let slot = self.command_id as usize % self.size;
1111        let mut cmd = command;
1112        unsafe {
1113            ptr::write(&mut cmd.command_id as *mut u16, self.command_id);
1114        }
1115        self.command_id = self.command_id.wrapping_add(1);
1116        self.submission.submit_command(cmd, slot);
1117        let deadline = tsc_deadline_ms(5500); // NVMe spec: up to 5.5s for admin commands
1118        loop {
1119            if let Some(c) = self.completion.poll_completion() {
1120                return Some(c);
1121            }
1122            if tsc_expired(deadline) {
1123                log::error!("NVMe: admin command timeout");
1124                return None;
1125            }
1126            core::hint::spin_loop();
1127        }
1128    }
1129}
1130
1131static NVME_CONTROLLERS: Mutex<Vec<Arc<Mutex<NvmeController>>>> = Mutex::new(Vec::new());
1132static NVME_INITIALIZED: AtomicBool = AtomicBool::new(false);
1133
1134pub static NVME_IRQ_LINE: AtomicU8 = AtomicU8::new(0);
1135
1136pub fn init() {
1137    log::info!("[NVMe] Scanning for NVMe controllers...");
1138
1139    let candidates = pci::probe_all(ProbeCriteria {
1140        vendor_id: None,
1141        device_id: None,
1142        class_code: Some(pci::class::MASS_STORAGE),
1143        subclass: Some(pci::storage_subclass::NVM),
1144        prog_if: None,
1145    });
1146
1147    for (i, pci_dev) in candidates.into_iter().enumerate() {
1148        log::info!(
1149            "NVMe: Found controller at {:?} (VEN:{:04x} DEV:{:04x})",
1150            pci_dev.address,
1151            pci_dev.vendor_id,
1152            pci_dev.device_id
1153        );
1154
1155        pci_dev.enable_bus_master();
1156        pci_dev.enable_memory_space();
1157
1158        // Try MSI/MSI-X first; fall back to INTx line.
1159        let (irq, vector) = crate::arch::x86_64::msi::probe_and_enable(&pci_dev, true);
1160
1161        let bar = match pci_dev.read_bar(0) {
1162            Some(Bar::Memory64 { addr, .. }) => addr,
1163            _ => {
1164                log::warn!("NVMe: Invalid BAR0");
1165                continue;
1166            }
1167        };
1168
1169        paging::ensure_identity_map_range(bar, 0x10000);
1170        let registers = phys_to_virt(bar) as usize;
1171        let name = format!("nvme{}", i);
1172
1173        match unsafe { NvmeController::new(registers, name.clone()) } {
1174            Ok(mut controller) => {
1175                controller.set_irq_line(irq);
1176                NVME_IRQ_LINE.store(vector, Ordering::Relaxed);
1177                log::info!(
1178                    "NVMe: {} initialized, IRQ={} vector={:#x}",
1179                    name,
1180                    irq,
1181                    vector
1182                );
1183                NVME_CONTROLLERS
1184                    .lock()
1185                    .push(Arc::new(Mutex::new(controller)));
1186                crate::arch::x86_64::idt::register_nvme_irq_vector(vector);
1187            }
1188            Err(e) => {
1189                log::warn!("NVMe: Failed to initialize controller: {:?}", e);
1190            }
1191        }
1192    }
1193
1194    NVME_INITIALIZED.store(true, Ordering::SeqCst);
1195    log::info!(
1196        "[NVMe] Found {} controller(s)",
1197        NVME_CONTROLLERS.lock().len()
1198    );
1199}
1200
1201pub fn get_first_controller() -> Option<Arc<Mutex<NvmeController>>> {
1202    NVME_CONTROLLERS.lock().first().cloned()
1203}
1204
1205pub fn is_available() -> bool {
1206    NVME_INITIALIZED.load(Ordering::Relaxed) && !NVME_CONTROLLERS.lock().is_empty()
1207}
1208
1209pub fn handle_interrupt() {
1210    if let Some(ctrl) = get_first_controller() {
1211        let controller = ctrl.lock();
1212        controller.handle_interrupt();
1213    }
1214}
1215
1216pub fn list_controllers() -> Vec<String> {
1217    NVME_CONTROLLERS
1218        .lock()
1219        .iter()
1220        .map(|c| c.lock().name.clone())
1221        .collect()
1222}
strat9_kernel/hardware/storage/nvme.rs

strat9_kernel/hardware/storage/
nvme.rs
