http-boot-client/src/multicast.rs

use alloc::boxed::Box;
use alloc::string::String;
use log::{debug, error, info};
use uefi::prelude::*;
use uefi::proto::network::pxe::{BaseCode, UdpOpFlags};
use uefi::table::boot::{ScopedProtocol, SearchType};
use uefi::{Error, Identify, Result};
extern crate alloc;
use crate::config::Config;
use uefi::data_types::CStr16;
use uefi::fs::{Path, PathBuf};
use uefi::proto::device_path::text::DevicePathFromText;
use uefi::proto::device_path::DevicePath;
use uefi::proto::network::IpAddress;

fn split_arg(raw_val: &str) -> Option<(&str, &str)> {
    let mut splitted = raw_val.split("=");
    let arg_name = match splitted.next() {
        None => return None,
        Some(v) => v,
    };
    let arg_value = match splitted.next() {
        None => return None,
        Some(v) => v,
    };
    return Some((arg_name, arg_value));
}

fn get_boot_path<'a>(
    boot_services: &BootServices,
    msg: String,
) -> Result<Box<DevicePath>, &'a str> {
    debug!("Initializing DevicePathFromTextProtocol");
    let device_path_from_text_handle = match boot_services
        .locate_handle_buffer(SearchType::ByProtocol(&DevicePathFromText::GUID))
    {
        Err(e) => {
            return Err(Error::new(e.status(), "failed to locate handle buffer"));
        }
        Ok(h) => match h.first() {
            None => return Err(Error::new(Status::UNSUPPORTED, "no handle buffer")),
            Some(h) => *h,
        },
    };
    let device_path_from_text = match boot_services
        .open_protocol_exclusive::<DevicePathFromText>(device_path_from_text_handle)
    {
        Err(e) => {
            return Err(Error::new(
                e.status(),
                "failed to open device_path_from_text protocol",
            ))
        }
        Ok(d) => d,
    };

    debug!("Splitting message");
    // a message has the following format: "ACTION arg1=val1,arg2=val2"
    let mut splitted_msg_iter = msg.split_ascii_whitespace();

    debug!("Checking if message is a valid BOOT_ACCEPT message");
    match splitted_msg_iter.next() {
        None => return Err(Error::new(Status::UNSUPPORTED, "empty message")),
        Some(action) => match action {
            "BOOT_ACCEPT" => (),
            "BOOT_DENY" => return Err(Error::new(Status::UNSUPPORTED, "boot refused")),
            _ => return Err(Error::new(Status::UNSUPPORTED, "unexpected action")),
        },
    }

    debug!("Parsing message arguments");
    match splitted_msg_iter.next() {
        None => return Err(Error::new(Status::UNSUPPORTED, "empty list of arguments")),
        Some(args_list) => {
            // split all the arguments in a comma separated list
            for arg in args_list.split(";") {
                match split_arg(arg) {
                    None => debug!("unexpected format for argument {}", arg),
                    Some((arg_name, arg_value)) => match arg_name {
                        "id" => debug!("found id {}", arg_value),
                        "efi_app" => {
                            debug!("found efi app {}", arg_value);
                            let mut buf = [0; 2000];
                            let path_raw = CStr16::from_str_with_buf(arg_value, &mut buf)
                                .expect("invalid str format");

                            return Ok(device_path_from_text
                                .convert_text_to_device_path(path_raw)
                                .expect("failed to parse device path")
                                .to_boxed());
                        }
                        _ => debug!("unexpected argument {} with value {}", arg_name, arg_value),
                    },
                }
            }
        }
    }
    Err(Error::new(Status::UNSUPPORTED, "no efi app defined"))
}

fn read_mcast_response<'a>(
    boot_services: &BootServices,
    base_client: &mut ScopedProtocol<'_, BaseCode>,
) -> Result<Box<DevicePath>> {
    let mut buf: [u8; 10000] = [0; 10000];

    let n = base_client
        .udp_read(
            UdpOpFlags::ANY_DEST_IP
                | UdpOpFlags::ANY_DEST_PORT
                | UdpOpFlags::ANY_SRC_IP
                | UdpOpFlags::ANY_SRC_PORT,
            None,
            None,
            None,
            None,
            None,
            &mut buf,
        )
        .unwrap();

    info!("Read {} bytes", n);
    let received_msg = String::from_utf8(buf[..n].to_vec()).unwrap();
    info!("Received message {}", received_msg);
    match get_boot_path(boot_services, received_msg) {
        Ok(r) => Ok(r),
        Err(e) => {
            error!("failed to get boot path: {}", e.data());
            Err(e.to_err_without_payload())
        }
    }
}

pub fn send_mcast<'a>(boot_services: &BootServices, conf: &Config) -> Result<Box<DevicePath>> {
    debug!("Initializing base bootservice client handle");
    let base_client_handle = *boot_services
        .locate_handle_buffer(SearchType::ByProtocol(&BaseCode::GUID))?
        .first()
        .expect("DevicePathToText is missing");

    let mut base_client = boot_services.open_protocol_exclusive::<BaseCode>(base_client_handle)?;

    let mut req_msg = String::from("BOOT_REQUEST id=");
    req_msg.push_str(&conf.client_uuid);

    info!("Starting PXE client");
    debug!("Starting PXE client");
    base_client.start(true)?;

    // debug!("Starting DHCP client");
    // base_client.dhcp(false)?;

    debug!("Setting IP address for PXE client");
    let mask = IpAddress::new_v6([
        0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00,
    ]);

    // 2001:abcd:1234:1::2/96
    let src_bytes: [u8; 16] = [
        0x20, 0x01, 0xab, 0xcd, 0x12, 0x34, 0x00, 0x01, 0x00, 0x0, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x02,
    ];
    let src_address: IpAddress = IpAddress::new_v6(src_bytes);
    // base_client
    base_client.set_station_ip(Some(&src_address), None)?;
    // base_client.set_station_ip(Some(&conf.source_ip), None)?;

    info!("Sending UDP multicast message");
    base_client.udp_write(
        UdpOpFlags::MAY_FRAGMENT,
        &conf.multicast_group,
        conf.multicast_port,
        None,
        // Some(&conf.source_ip),
        Some(&src_address),
        // None,
        None,
        None,
        req_msg.as_bytes(),
    )?;
    let res = read_mcast_response(boot_services, &mut base_client);

    info!("Stopping PXE client");
    base_client.stop()?;

    info!("PXE client successfully executed");
    res
}