yggdrasil/kernel/driver/bus/usb/src/pipe/control.rs

use core::{
    mem::{size_of, MaybeUninit},
    ops::Deref,
};

use alloc::{boxed::Box, string::String};
use async_trait::async_trait;
use bytemuck::{Pod, Zeroable};
use libk_mm::PageBox;

use crate::{
    descriptor::{
        UsbConfigurationDescriptor, UsbDeviceDescriptor, UsbDeviceQualifier, UsbEndpointDescriptor,
        UsbInterfaceDescriptor, UsbOtherSpeedConfiguration,
    },
    error::UsbError,
};

#[derive(Debug)]
pub struct ControlTransferSetup {
    pub bm_request_type: u8,
    pub b_request: u8,
    pub w_value: u16,
    pub w_index: u16,
    pub w_length: u16,
}

#[derive(Clone, Copy, Debug, Default, Pod, Zeroable)]
#[repr(C)]
pub struct SetConfiguration;

pub trait UsbDeviceRequest: Sized + Pod {
    const BM_REQUEST_TYPE: u8;
    const B_REQUEST: u8;
}

pub trait UsbClassSpecificRequest: Sized + Pod {
    const BM_REQUEST_TYPE: u8;
    const B_REQUEST: u8;
}

pub trait UsbDescriptorRequest: UsbDeviceRequest {
    const DESCRIPTOR_TYPE: u8;
}

impl UsbDescriptorRequest for UsbDeviceDescriptor {
    const DESCRIPTOR_TYPE: u8 = 1;
}

impl UsbDescriptorRequest for UsbConfigurationDescriptor {
    const DESCRIPTOR_TYPE: u8 = 2;
}

impl UsbDescriptorRequest for UsbInterfaceDescriptor {
    const DESCRIPTOR_TYPE: u8 = 4;
}

impl UsbDeviceRequest for SetConfiguration {
    const BM_REQUEST_TYPE: u8 = 0;
    const B_REQUEST: u8 = 0x09;
}

impl<U: UsbDescriptorRequest> UsbDeviceRequest for U {
    const BM_REQUEST_TYPE: u8 = 0b10000000;
    const B_REQUEST: u8 = 0x06;
}

fn decode_usb_string(bytes: &[u8]) -> Result<String, UsbError> {
    if bytes.len() % 2 != 0 {
        return Err(UsbError::InvalidDescriptorField);
    }

    char::decode_utf16(
        bytes
            .iter()
            .array_chunks::<2>()
            .map(|[&a, &b]| u16::from_le_bytes([a, b])),
    )
    .collect::<Result<String, _>>()
    .map_err(|_| UsbError::InvalidDescriptorField)
}

// Pipe impl

#[async_trait]
pub trait UsbControlPipe: Send + Sync {
    async fn control_transfer(&self, setup: ControlTransferSetup) -> Result<(), UsbError>;
    async fn control_transfer_in(
        &self,
        setup: ControlTransferSetup,
        buffer: &mut [MaybeUninit<u8>],
    ) -> Result<usize, UsbError>;
    async fn control_transfer_out(
        &self,
        setup: ControlTransferSetup,
        buffer: &[u8],
    ) -> Result<usize, UsbError>;
}

pub struct UsbControlPipeAccess(pub Box<dyn UsbControlPipe>);

#[derive(Debug)]
pub enum ConfigurationDescriptorEntry<'a> {
    Configuration(&'a UsbConfigurationDescriptor),
    Interface(&'a UsbInterfaceDescriptor),
    Endpoint(&'a UsbEndpointDescriptor),
    DeviceQualifier(&'a UsbDeviceQualifier),
    OtherSpeed(&'a UsbOtherSpeedConfiguration),
    Other,
}

pub struct ConfigurationDescriptorIter<'a> {
    buffer: &'a PageBox<[u8]>,
    offset: usize,
}

pub struct ConfigurationDescriptorQuery {
    buffer: PageBox<[u8]>,
}

impl<'a> Iterator for ConfigurationDescriptorIter<'a> {
    type Item = ConfigurationDescriptorEntry<'a>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.offset + 2 >= self.buffer.len() {
            return None;
        }

        let desc_len = self.buffer[self.offset] as usize;
        let desc_ty = self.buffer[self.offset + 1];

        if desc_len == 0 {
            return None;
        }

        let entry = match desc_ty {
            0x02 if desc_len == size_of::<UsbConfigurationDescriptor>() => {
                ConfigurationDescriptorEntry::Configuration(bytemuck::from_bytes(
                    &self.buffer[self.offset..self.offset + desc_len],
                ))
            }
            0x04 if desc_len == size_of::<UsbInterfaceDescriptor>() => {
                ConfigurationDescriptorEntry::Interface(bytemuck::from_bytes(
                    &self.buffer[self.offset..self.offset + desc_len],
                ))
            }
            0x05 if desc_len == size_of::<UsbEndpointDescriptor>() => {
                ConfigurationDescriptorEntry::Endpoint(bytemuck::from_bytes(
                    &self.buffer[self.offset..self.offset + desc_len],
                ))
            }
            0x07 if desc_len == size_of::<UsbOtherSpeedConfiguration>() => {
                ConfigurationDescriptorEntry::OtherSpeed(bytemuck::from_bytes(
                    &self.buffer[self.offset..self.offset + desc_len],
                ))
            }
            _ => ConfigurationDescriptorEntry::Other,
        };

        self.offset += desc_len;

        Some(entry)
    }
}

impl ConfigurationDescriptorQuery {
    pub fn configuration(&self) -> &UsbConfigurationDescriptor {
        bytemuck::from_bytes(&self.buffer[..size_of::<UsbConfigurationDescriptor>()])
    }

    pub fn descriptors(&self) -> ConfigurationDescriptorIter<'_> {
        ConfigurationDescriptorIter {
            buffer: &self.buffer,
            offset: 0,
        }
    }
}

impl UsbControlPipeAccess {
    pub async fn query_device_descriptor(&self) -> Result<UsbDeviceDescriptor, UsbError> {
        let mut buffer = MaybeUninit::uninit();

        self.control_transfer_in(
            ControlTransferSetup {
                bm_request_type: 0b10000000,
                b_request: 0x06,
                w_value: 0x100,
                w_index: 0,
                w_length: size_of::<UsbDeviceDescriptor>() as _,
            },
            buffer.as_bytes_mut(),
        )
        .await?;

        Ok(unsafe { buffer.assume_init() })
    }

    async fn fill_configuation_descriptor(
        &self,
        index: u8,
        buffer: &mut [MaybeUninit<u8>],
    ) -> Result<(), UsbError> {
        self.control_transfer_in(
            ControlTransferSetup {
                bm_request_type: 0b10000000,
                b_request: 0x06,
                w_value: 0x200 | (index as u16),
                w_index: 0,
                w_length: buffer.len().try_into().unwrap(),
            },
            buffer,
        )
        .await?;
        Ok(())
    }

    pub async fn query_string(&self, index: u8) -> Result<String, UsbError> {
        let mut buffer = MaybeUninit::uninit_array::<256>();

        let len = self
            .control_transfer_in(
                ControlTransferSetup {
                    bm_request_type: 0b10000000,
                    b_request: 0x06,
                    w_value: 0x300 | (index as u16),
                    w_index: 0,
                    w_length: 4096,
                },
                &mut buffer[..],
            )
            .await?;
        let data = unsafe { MaybeUninit::slice_assume_init_ref(&buffer[..len]) };
        let len = data[0] as usize;

        decode_usb_string(&data[2..len])
    }

    pub async fn query_configuration_descriptor(
        &self,
        index: u8,
    ) -> Result<ConfigurationDescriptorQuery, UsbError> {
        // 4K should be enough for a configuration descriptor
        let mut buffer = PageBox::new_uninit_slice(4096).map_err(UsbError::MemoryError)?;

        self.fill_configuation_descriptor(index, &mut buffer[..])
            .await?;

        let buffer = unsafe { PageBox::assume_init_slice(buffer) };

        let desc: &UsbConfigurationDescriptor =
            bytemuck::from_bytes(&buffer[..size_of::<UsbConfigurationDescriptor>()]);
        let total_len = desc.total_length as usize;

        if total_len > 4096 {
            unimplemented!("4KiB wasn't enough");
        }

        Ok(ConfigurationDescriptorQuery { buffer })
    }

    pub async fn perform_action<D: UsbDeviceRequest>(
        &self,
        w_value: u16,
        w_index: u16,
    ) -> Result<(), UsbError> {
        self.control_transfer(ControlTransferSetup {
            bm_request_type: D::BM_REQUEST_TYPE,
            b_request: D::B_REQUEST,
            w_value,
            w_index,
            w_length: 0,
        })
        .await
    }

    pub async fn class_specific_request<D: UsbClassSpecificRequest>(
        &self,
        w_value: u16,
        w_index: u16,
    ) -> Result<(), UsbError> {
        self.control_transfer(ControlTransferSetup {
            bm_request_type: D::BM_REQUEST_TYPE,
            b_request: D::B_REQUEST,
            w_value,
            w_index,
            w_length: 0,
        })
        .await
    }

    pub async fn set_configuration(&self, value: u16) -> Result<(), UsbError> {
        self.perform_action::<SetConfiguration>(value, 0).await
    }
}

impl Deref for UsbControlPipeAccess {
    type Target = dyn UsbControlPipe;

    fn deref(&self) -> &Self::Target {
        &*self.0
    }
}