RedBear-OS/drivers/usb/usbhubd/src/main.rs

use std::{env, thread, time};

use xhcid_interface::{
    plain, usb, ConfigureEndpointsReq, DevDesc, DeviceReqData, EndpointTy, EndpDesc,
    EndpDirection, PortId, PortReqRecipient, PortReqTy, XhciClientHandle, XhciEndpHandle,
};

fn main() {
    common::init();
    let mut args = env::args().skip(1);

    const USAGE: &'static str = "usbhubd <scheme> <port> <interface>";

    let scheme = args.next().expect(USAGE);
    let port_id = args
        .next()
        .expect(USAGE)
        .parse::<PortId>()
        .expect("Expected port ID");
    let interface_num = args
        .next()
        .expect(USAGE)
        .parse::<u8>()
        .expect("Expected integer as input of interface");

    log::info!(
        "USB HUB driver spawned with scheme `{}`, port {}, interface {}",
        scheme,
        port_id,
        interface_num
    );

    let name = format!("{}_{}_{}_hub", scheme, port_id, interface_num);
    common::setup_logging(
        "usb",
        "device",
        &name,
        common::output_level(),
        common::file_level(),
    );

    let handle =
        XhciClientHandle::new(scheme.clone(), port_id).expect("Failed to open XhciClientHandle");
    let desc: DevDesc = handle
        .get_standard_descs()
        .expect("Failed to get standard descriptors");

    let (conf_desc, if_desc) = desc
        .config_descs
        .iter()
        .find_map(|conf_desc| {
            let if_desc = conf_desc.interface_descs.iter().find_map(|if_desc| {
                if if_desc.number == interface_num {
                    Some(if_desc.clone())
                } else {
                    None
                }
            })?;
            Some((conf_desc.clone(), if_desc))
        })
        .expect("Failed to find suitable configuration");

    // Read hub descriptor
    let (ports, usb_3, b_pwr_on_2_pwr_good, w_hub_delay) = if desc.major_version() >= 3 {
        let mut hub_desc = usb::HubDescriptorV3::default();
        handle
            .device_request(
                PortReqTy::Class,
                PortReqRecipient::Device,
                usb::SetupReq::GetDescriptor as u8,
                u16::from(usb::HubDescriptorV3::DESCRIPTOR_KIND) << 8,
                0,
                DeviceReqData::In(unsafe { plain::as_mut_bytes(&mut hub_desc) }),
            )
            .expect("Failed to read hub descriptor");
        (
            hub_desc.ports,
            true,
            // USB 3 hub descriptor does not carry bPwrOn2PwrGood;
            // the spec says to use 10 (20ms) as default.
            10u8,
            u16::from(hub_desc.delay),
        )
    } else {
        let mut hub_desc = usb::HubDescriptorV2::default();
        handle
            .device_request(
                PortReqTy::Class,
                PortReqRecipient::Device,
                usb::SetupReq::GetDescriptor as u8,
                u16::from(usb::HubDescriptorV2::DESCRIPTOR_KIND) << 8,
                0,
                DeviceReqData::In(unsafe { plain::as_mut_bytes(&mut hub_desc) }),
            )
            .expect("Failed to read hub descriptor");
        (
            hub_desc.ports,
            false,
            hub_desc.power_on_good,
            0u16, // wHubDelay only exists in USB 3 hub descriptors
        )
    };

    log::info!(
        "usbhubd: {} port(s) detected, PwrOn2PwrGood={}*2ms, wHubDelay={}",
        ports, b_pwr_on_2_pwr_good, w_hub_delay,
    );

    // Configure as hub device.  USB 3 hubs do not need explicit
    // interface_desc / alternate_setting in ConfigureEndpointsReq —
    // the xHCI controller derives those from the default alternate
    // (alt 0).  Passing Some(...) causes stalls on USB 3 hubs.
    handle
        .configure_endpoints(&ConfigureEndpointsReq {
            config_desc: conf_desc.configuration_value,
            interface_desc: if usb_3 {
                None
            } else {
                Some(interface_num)
            },
            alternate_setting: if usb_3 {
                None
            } else {
                Some(if_desc.alternate_setting)
            },
            hub_ports: Some(ports),
        })
        .expect("Failed to configure endpoints after reading hub descriptor");

    // SET_HUB_DEPTH for USB 3 hubs (Linux 7.1 hub_activate).
    if usb_3 {
        let hub_depth = port_id.hub_depth();
        handle
            .device_request(
                PortReqTy::Class,
                PortReqRecipient::Device,
                0x0c, // SET_HUB_DEPTH
                u16::from(hub_depth),
                0,
                DeviceReqData::NoData,
            )
            .expect("Failed to set hub depth");
        log::info!("usbhubd: SET_HUB_DEPTH to {}", hub_depth);
    }

    // Find the hub interrupt endpoint for status-change detection.
    // Linux 7.1 hub_irq() reads the status-change bitmap from EP1.
    // For USB 2.x hubs: EP1, bInterval=12 (255ms max polling).
    // For USB 3 hubs: may not have a dedicated interrupt EP — fall back to polling.
    let intr_desc = if_desc
        .endpoints
        .iter()
        .find(|ep| ep.ty() == EndpointTy::Interrupt && ep.direction() == EndpDirection::In)
        .cloned();

    let mut intr_ep_handle: Option<XhciEndpHandle> = if intr_desc.is_some() {
        match handle.open_endpoint(1u8) {
            Ok(handle) => {
                log::info!("usbhubd: interrupt endpoint opened for change detection");
                Some(handle)
            }
            Err(e) => {
                log::warn!("usbhubd: interrupt endpoint open failed ({}), falling back to polling", e);
                None
            }
        }
    } else {
        log::info!("usbhubd: no interrupt endpoint found, using polling");
        None
    };

    // Initialize states
    struct PortState {
        port_id: PortId,
        port_sts: usb::HubPortStatus,
        handle: XhciClientHandle,
        attached: bool,
    }

    impl PortState {
        pub fn ensure_attached(&mut self, attached: bool) {
            if attached == self.attached {
                return;
            }

            if attached {
                match self.handle.attach() {
                    Ok(()) => {}
                    Err(e) => {
                        log::warn!("usbhubd: attach failed for port: {}", e);
                        return;
                    }
                }
            } else {
                match self.handle.detach() {
                    Ok(()) => {}
                    Err(e) => {
                        log::warn!("usbhubd: detach failed for port: {}", e);
                        return;
                    }
                }
            }

            self.attached = attached;
        }
    }

    let mut states = Vec::new();
    for port in 1..=ports {
        let child_port_id = port_id.child(port).expect("Cannot get child port ID");
        states.push(PortState {
            port_id: child_port_id,
            port_sts: if usb_3 {
                usb::HubPortStatus::V3(usb::HubPortStatusV3::default())
            } else {
                usb::HubPortStatus::V2(usb::HubPortStatusV2::default())
            },
            handle: XhciClientHandle::new(scheme.clone(), child_port_id)
                .expect("Failed to open XhciClientHandle"),
            attached: false,
        });
    }

    // Power-on delay in milliseconds per Linux 7.1 hub_power_on_good_delay.
    let power_on_delay_ms = u64::from(b_pwr_on_2_pwr_good) * 2;
    // Linux uses 100ms minimum; follow the same floor.
    let power_on_delay_ms = core::cmp::max(power_on_delay_ms, 100);

    // Main event loop with interrupt-driven change detection.
    // Linux 7.1 hub_irq(): reads status-change bitmap from EP1,
    // then kicks hub_wq to process only the changed ports.
    // We do the same: read the bitmap, build a port mask, and
    // only poll GetPortStatus for ports whose bit is set.
    // If EP1 is unavailable, fall back to polling all ports.
    //
    // Bitmap size: ceil(ports / 8) bytes.
    // Port N is bit (N-1) of byte (N-1)/8.
    const POLL_FALLBACK_MS: u64 = 200;
    let bitmap_size = (ports as usize + 7) / 8;
    let mut bitmap = vec![0u8; bitmap_size];

    loop {
        // Build a port-change mask.
        // Bit N set = port (N+1) needs processing.
        let changed: u64 = if let Some(ref mut ep) = intr_ep_handle {
            match ep.transfer_read(&mut bitmap) {
                Ok(_) => {
                    let mut mask = 0u64;
                    for (byte_idx, &byte) in bitmap.iter().enumerate() {
                        if byte != 0 {
                            mask |= (byte as u64) << (byte_idx * 8);
                        }
                    }
                    mask
                }
                Err(e) => {
                    log::warn!("usbhubd: interrupt transfer failed ({}), falling back to poll", e);
                    (1u64 << ports) - 1
                }
            }
        } else {
            // Polling mode: process all ports.
            thread::sleep(time::Duration::from_millis(POLL_FALLBACK_MS));
            (1u64 << ports) - 1
        };
        for port in 1..=ports {
            let bit = 1u64 << (port - 1);
            if (changed & bit) == 0 {
                continue;
            }

            let port_idx: usize = match port.checked_sub(1) {
                Some(p) => p.into(),
                None => continue,
            };
            let state = match states.get_mut(port_idx) {
                Some(s) => s,
                None => continue,
            };

            let port_sts = if usb_3 {
                let mut port_sts = usb::HubPortStatusV3::default();
                match handle.device_request(
                    PortReqTy::Class,
                    PortReqRecipient::Other,
                    usb::SetupReq::GetStatus as u8,
                    0,
                    port as u16,
                    DeviceReqData::In(unsafe { plain::as_mut_bytes(&mut port_sts) }),
                ) {
                    Ok(()) => usb::HubPortStatus::V3(port_sts),
                    Err(e) => {
                        log::warn!("usbhubd: GetPortStatus failed for port {}: {} — detaching", port, e);
                        state.ensure_attached(false);
                        continue;
                    }
                }
            } else {
                let mut port_sts = usb::HubPortStatusV2::default();
                match handle.device_request(
                    PortReqTy::Class,
                    PortReqRecipient::Other,
                    usb::SetupReq::GetStatus as u8,
                    0,
                    port as u16,
                    DeviceReqData::In(unsafe { plain::as_mut_bytes(&mut port_sts) }),
                ) {
                    Ok(()) => usb::HubPortStatus::V2(port_sts),
                    Err(e) => {
                        log::warn!("usbhubd: GetPortStatus failed for port {}: {} — detaching", port, e);
                        state.ensure_attached(false);
                        continue;
                    }
                }
            };
            if state.port_sts != port_sts {
                state.port_sts = port_sts;
                log::info!("port {} status {:X?}", port, port_sts);
            }

            // Ensure port is powered on.
            // Linux 7.1 hub_power_on(): issue SET_FEATURE(PORT_POWER),
            // then sleep bPwrOn2PwrGood * 2ms (minimum 100ms).
            if !port_sts.is_powered() {
                log::info!("power on port {port}");
                if let Err(e) = handle.device_request(
                    PortReqTy::Class,
                    PortReqRecipient::Other,
                    usb::SetupReq::SetFeature as u8,
                    usb::HubPortFeature::PortPower as u16,
                    port as u16,
                    DeviceReqData::NoData,
                ) {
                    log::warn!("usbhubd: SetPortPower failed for port {}: {}", port, e);
                    continue;
                }
                state.ensure_attached(false);
                thread::sleep(time::Duration::from_millis(power_on_delay_ms));
                continue;
            }

            // Linux 7.1: over-current detection and recovery.
            // C_PORT_OVERCURRENT → log, clear, power-cycle.
            if port_sts.is_over_current_changed() {
                log::warn!("usbhubd: over-current change on port {} — power cycling", port);
                let _ = handle.device_request(
                    PortReqTy::Class,
                    PortReqRecipient::Other,
                    usb::SetupReq::ClearFeature as u8,
                    usb::HubPortFeature::CPortOverCurrent as u16,
                    port as u16,
                    DeviceReqData::NoData,
                );
                let _ = handle.device_request(
                    PortReqTy::Class,
                    PortReqRecipient::Other,
                    usb::SetupReq::ClearFeature as u8,
                    usb::HubPortFeature::PortPower as u16,
                    port as u16,
                    DeviceReqData::NoData,
                );
                state.ensure_attached(false);
                thread::sleep(time::Duration::from_millis(power_on_delay_ms));
                continue;
            }

            // Linux 7.1: port indicator — set amber during reset/power-on,
            // green when enabled.  Helps diagnose which port is active.
            if port_sts.is_connected() && port_sts.is_enabled() && !port_sts.is_resetting() {
                let _ = handle.device_request(
                    PortReqTy::Class,
                    PortReqRecipient::Other,
                    usb::SetupReq::SetFeature as u8,
                    usb::HubPortFeature::PortIndicator as u16,
                    port as u16,
                    DeviceReqData::NoData,
                );
            }

            // Ignore disconnected port
            if !port_sts.is_connected() {
                state.ensure_attached(false);
                continue;
            }

            // Ignore port in reset
            if port_sts.is_resetting() {
                state.ensure_attached(false);
                continue;
            }

            // Ensure port is enabled.
            // Linux 7.1 hub_port_reset(): issue SET_FEATURE(PORT_RESET),
            // then wait for reset completion (up to USB_PORT_RESET_TIMEOUT
            // = 5000ms for USB 3, 1000ms for USB 2).
            if !port_sts.is_enabled() {
                log::info!("reset port {port}");
                if let Err(e) = handle.device_request(
                    PortReqTy::Class,
                    PortReqRecipient::Other,
                    usb::SetupReq::SetFeature as u8,
                    usb::HubPortFeature::PortReset as u16,
                    port as u16,
                    DeviceReqData::NoData,
                ) {
                    log::warn!("usbhubd: SetPortReset failed for port {}: {}", port, e);
                    continue;
                }
                state.ensure_attached(false);
                thread::sleep(time::Duration::from_millis(10));
                continue;
            }

            state.ensure_attached(true);
        }
    }
}