yggdrasil/kernel/driver/block/ahci/src/lib.rs

#![feature(generic_const_exprs)]
#![allow(incomplete_features)]
#![no_std]

extern crate alloc;

use alloc::{format, sync::Arc, vec::Vec};
use bytemuck::Zeroable;
use data::ReceivedFis;
use device_api::{
    device::Device,
    interrupt::{InterruptAffinity, InterruptHandler},
};
use error::AhciError;
use libk::{device::manager::probe_partitions, fs::devfs, task::runtime};
use libk_mm::{address::AsPhysicalAddress, device::DeviceMemoryIo, PageBox};
use libk_util::{sync::IrqSafeSpinlock, OneTimeInit};
use port::AhciPort;
use regs::{PortRegs, Regs};
use tock_registers::interfaces::{ReadWriteable, Readable, Writeable};
use ygg_driver_pci::{
    device::{PciDeviceInfo, PreferredInterruptMode},
    PciCommandRegister, PciConfigurationSpace,
};
use yggdrasil_abi::{error::Error, io::FileMode};

use crate::regs::{Version, CAP, GHC, SSTS};

mod command;
mod data;
mod error;
mod port;
mod regs;

const MAX_PRD_SIZE: usize = 65536;
const MAX_COMMANDS: usize = u32::BITS as usize;
const SECTOR_SIZE: usize = 512;
const MAX_DRIVES: usize = (b'z' - b'a') as usize;

pub struct AhciController {
    regs: IrqSafeSpinlock<DeviceMemoryIo<'static, Regs>>,
    ports: OneTimeInit<Vec<Arc<AhciPort>>>,
    received_fis_buffers: OneTimeInit<[Option<PageBox<ReceivedFis>>; 16]>,

    version: Version,
    max_port_count: usize,
    ahci_only: bool,
    has_64_bit: bool,
}

impl AhciController {
    async fn late_init(self: Arc<Self>) -> Result<(), AhciError> {
        log::info!("Initializing AHCI SATA Controller {:?}", self.version);

        let regs = self.regs.lock();

        regs.GHC.modify(GHC::HR::SET);

        while regs.GHC.matches_all(GHC::HR::SET) {
            core::hint::spin_loop();
        }

        if !self.ahci_only {
            regs.GHC.modify(GHC::AE::SET);
        }

        let pi = regs.PI.get();

        let mut ports = Vec::new();

        drop(regs);

        let mut fis_buffers = [const { None }; 16];
        // Allocate FIS receive buffers for the ports
        for (i, fis_buffer_slot) in fis_buffers.iter_mut().enumerate().take(self.max_port_count) {
            if pi & (1 << i) == 0 {
                continue;
            }

            let regs = self.regs.lock();
            let port = &regs.PORTS[i];

            let buffer = PageBox::new(ReceivedFis::zeroed()).map_err(AhciError::MemoryError)?;
            port.set_received_fis_address_64(unsafe { buffer.as_physical_address() });
            *fis_buffer_slot = Some(buffer);
        }

        self.received_fis_buffers.init(fis_buffers);

        for i in 0..self.max_port_count {
            if pi & (1 << i) == 0 {
                continue;
            }

            let regs = self.regs.lock();
            let port = &regs.PORTS[i];

            if !port.SSTS.matches_all(SSTS::DET::Online + SSTS::IPM::Active) {
                continue;
            }

            port.start()?;

            // TODO wait here

            let sig = port.SIG.get();
            if sig != PortRegs::SIG_SATA {
                log::warn!("Skipping unknown port {} with signature {:#x}", i, sig);
                continue;
            }

            let port = unsafe { regs.extract(|regs| &regs.PORTS[i]) };

            drop(regs);

            let port = match AhciPort::create(port, self.clone(), i) {
                Ok(port) => port,
                Err(error) => {
                    log::warn!("Port {} init error: {:?}", i, error);
                    continue;
                }
            };

            ports.push(port);
        }

        let ports = self.ports.init(ports);

        // Enable global HC interrupts
        self.regs.lock().GHC.modify(GHC::IE::SET);

        // Setup the detected ports
        for (i, port) in ports.iter().enumerate() {
            log::info!("Init port {}", i);
            port.init_inner().await?;
        }

        // Dump info about the drives
        for (i, port) in ports.iter().enumerate() {
            let info = port.info().unwrap();
            log::info!(
                "Port {}: model={:?}, serial={:?}, lba_count={}",
                i,
                info.model,
                info.serial,
                info.lba_count
            );
        }

        for port in ports.iter() {
            register_sata_drive(port.clone(), true);
        }

        log::debug!("All ports initialized");

        Ok(())
    }
}

impl InterruptHandler for AhciController {
    fn handle_irq(self: Arc<Self>, _vector: Option<usize>) -> bool {
        let regs = self.regs.lock();

        let is = regs.IS.get();
        if is != 0 {
            if let Some(ports) = self.ports.try_get() {
                // Clear global interrupt status
                regs.IS.set(u32::MAX);

                for port in ports {
                    if is & (1 << port.index) != 0 {
                        port.handle_pending_interrupts();
                    }
                }
            }
        }

        false
    }
}

impl Device for AhciController {
    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
        // Do the init in background
        runtime::spawn(self.late_init())?;
        Ok(())
    }

    fn display_name(&self) -> &str {
        "AHCI Controller"
    }
}

static SATA_DRIVES: IrqSafeSpinlock<Vec<Arc<AhciPort>>> = IrqSafeSpinlock::new(Vec::new());

pub fn probe(info: &PciDeviceInfo) -> Result<Arc<dyn Device>, Error> {
    let bar5 = info.config_space.bar(5).ok_or(Error::InvalidOperation)?;
    let bar5 = bar5.as_memory().ok_or(Error::InvalidOperation)?;

    let mut cmd = PciCommandRegister::from_bits_retain(info.config_space.command());
    cmd &= !(PciCommandRegister::DISABLE_INTERRUPTS | PciCommandRegister::ENABLE_IO);
    cmd |= PciCommandRegister::ENABLE_MEMORY | PciCommandRegister::BUS_MASTER;
    info.config_space.set_command(cmd.bits());

    info.init_interrupts(PreferredInterruptMode::Msi(true))?;

    // // TODO support regular PCI interrupts (ACPI dependency)
    // let Some(mut msi) = info.config_space.capability::<MsiCapability>() else {
    //     log::warn!("Ignoring AHCI: does not support MSI (and the OS doesn't yet support PCI IRQ)");
    //     return Err(Error::InvalidOperation);
    // };

    // Map the registers
    let regs = unsafe { DeviceMemoryIo::<Regs>::map(bar5, Default::default()) }?;

    let version = Version::try_from(regs.VS.get())?;
    let ahci_only = regs.CAP.matches_all(CAP::SAM::SET);
    let max_port_count = regs.CAP.read(CAP::NP) as usize;
    let has_64_bit = regs.CAP.matches_all(CAP::S64A::SET);

    // TODO extract Number of Command Slots

    let ahci = Arc::new(AhciController {
        regs: IrqSafeSpinlock::new(regs),
        ports: OneTimeInit::new(),
        received_fis_buffers: OneTimeInit::new(),
        version,
        max_port_count,
        ahci_only,
        has_64_bit,
    });

    // TODO use multiple vectors if capable
    info.map_interrupt(InterruptAffinity::Any, ahci.clone())?;

    Ok(ahci)
}

pub fn register_sata_drive(drive: Arc<AhciPort>, probe: bool) {
    let index = {
        let mut drives = SATA_DRIVES.lock();
        let index = drives.len();
        if index >= MAX_DRIVES {
            log::error!("Cannot add a SATA drive: too many of them");
            return;
        }
        drives.push(drive.clone());
        index
    };
    let letter = (index as u8 + b'a') as char;

    let name = format!("sd{letter}");
    log::info!("Register SATA drive: {name}");

    devfs::add_named_block_device(drive.clone(), name.clone(), FileMode::new(0o600)).ok();

    if probe {
        runtime::spawn(async move {
            let name = name;
            log::info!("Probing partitions for {name}");
            probe_partitions(drive, |index, partition| {
                let partition_name = format!("{name}{}", index + 1);
                devfs::add_named_block_device(
                    Arc::new(partition),
                    partition_name,
                    FileMode::new(0o600),
                )
                .ok();
            })
            .await
            .ok();
        })
        .ok();
    }
}