Initial Commit

src/main.rs

@@ -0,0 +1,185 @@
+mod utils;
+
+use crate::utils::{elf_to_bin, flash_bin, vendor_map};
+use colored::*;
+use rusb::{open_device_with_vid_pid, GlobalContext};
+
+use std::path::PathBuf;
+use std::process::{Command, Stdio};
+use std::time::Instant;
+use structopt::StructOpt;
+
+fn main() {
+    // Initialize the logging backend.
+    pretty_env_logger::init();
+
+    // Get commandline options.
+    // Skip the first arg which is the calling application name.
+    let opt = Opt::from_iter(std::env::args().skip(1));
+
+    // Try and get the cargo project information.
+    let project = cargo_project::Project::query(".").expect("Couldn't parse the Cargo.toml");
+
+    // Decide what artifact to use.
+    let artifact = if let Some(bin) = &opt.bin {
+        cargo_project::Artifact::Bin(bin)
+    } else if let Some(example) = &opt.example {
+        cargo_project::Artifact::Example(example)
+    } else {
+        cargo_project::Artifact::Bin(project.name())
+    };
+
+    // Decide what profile to use.
+    let profile = if opt.release {
+        cargo_project::Profile::Release
+    } else {
+        cargo_project::Profile::Dev
+    };
+
+    // Try and get the artifact path.
+    let path = project
+        .path(
+            artifact,
+            profile,
+            opt.target.as_deref(),
+            "x86_64-unknown-linux-gnu",
+        )
+        .expect("Couldn't find the build result");
+
+    // Remove first two args which is the calling application name and the `dfu` command from cargo.
+    let mut args: Vec<_> = std::env::args().skip(2).collect();
+
+    // todo, keep as iter. difficult because we want to filter map remove two items at once.
+    // Remove our args as cargo build does not understand them.
+    let flags = ["--pid", "--vid"].iter();
+    for flag in flags {
+        if let Some(index) = args.iter().position(|x| x == flag) {
+            args.remove(index);
+            args.remove(index);
+        }
+    }
+
+    let status = Command::new("cargo")
+        .arg("build")
+        .args(args)
+        .stdout(Stdio::inherit())
+        .stderr(Stdio::inherit())
+        .spawn()
+        .unwrap()
+        .wait()
+        .unwrap();
+
+    if !status.success() {
+        exit_with_process_status(status)
+    }
+
+    let d = if let (Some(v), Some(p)) = (opt.vid, opt.pid) {
+        open_device_with_vid_pid(v, p)
+            .expect("Are you sure device is plugged in and in bootloader mode?")
+    } else {
+        println!(
+            "    {} for a connected device with known vid/pid pair.",
+            "Searching".green().bold(),
+        );
+
+        let devices: Vec<_> = rusb::devices()
+            .expect("Error with Libusb")
+            .iter()
+            .map(|d| d.device_descriptor().unwrap())
+            .collect();
+
+        let mut device: Option<rusb::DeviceHandle<GlobalContext>> = None;
+
+        let vendor = vendor_map();
+
+        for d in devices {
+            if let Some(products) = vendor.get(&d.vendor_id()) {
+                if products.contains(&d.product_id()) {
+                    if let Some(d) = open_device_with_vid_pid(d.vendor_id(), d.product_id()) {
+                        device = Some(d);
+                        break;
+                    }
+                }
+            }
+        }
+        device.expect("Are you sure device is plugged in and in bootloader mode?")
+    };
+
+    println!(
+        "    {} {} {}",
+        "Trying ".green().bold(),
+        d.read_manufacturer_string_ascii(&d.device().device_descriptor().unwrap())
+            .unwrap(),
+        d.read_product_string_ascii(&d.device().device_descriptor().unwrap())
+            .unwrap()
+    );
+
+    println!("    {} {:?}", "Flashing".green().bold(), path);
+
+    let (binary, address) = elf_to_bin(path).unwrap();
+
+    // Start timer.
+    let instant = Instant::now();
+
+    // let bininfo = hf2::bin_info(&d).expect("bin_info failed");
+    // log::debug!("{:?}", bininfo);
+    flash_bin(&binary, address, &d.device()).unwrap();
+
+    // Stop timer.
+    let elapsed = instant.elapsed();
+    println!(
+        "    {} in {}s",
+        "Finished".green().bold(),
+        elapsed.as_millis() as f32 / 1000.0
+    );
+}
+
+#[cfg(unix)]
+fn exit_with_process_status(status: std::process::ExitStatus) -> ! {
+    use std::os::unix::process::ExitStatusExt;
+    let status = status.code().or_else(|| status.signal()).unwrap_or(1);
+    std::process::exit(status)
+}
+
+#[cfg(not(unix))]
+fn exit_with_process_status(status: std::process::ExitStatus) -> ! {
+    let status = status.code().unwrap_or(1);
+    std::process::exit(status)
+}
+
+fn parse_hex_16(input: &str) -> Result<u16, std::num::ParseIntError> {
+    if let Some(stripped) = input.strip_prefix("0x") {
+        u16::from_str_radix(stripped, 16)
+    } else {
+        input.parse::<u16>()
+    }
+}
+
+#[allow(unused)]
+#[derive(Debug, StructOpt)]
+struct Opt {
+    // `cargo build` arguments
+    #[structopt(name = "binary", long = "bin")]
+    bin: Option<String>,
+    #[structopt(name = "example", long = "example")]
+    example: Option<String>,
+    #[structopt(name = "package", short = "p", long = "package")]
+    package: Option<String>,
+    #[structopt(name = "release", long = "release")]
+    release: bool,
+    #[structopt(name = "target", long = "target")]
+    target: Option<String>,
+    #[structopt(name = "PATH", long = "manifest-path", parse(from_os_str))]
+    manifest_path: Option<PathBuf>,
+    #[structopt(long)]
+    no_default_features: bool,
+    #[structopt(long)]
+    all_features: bool,
+    #[structopt(long)]
+    features: Vec<String>,
+
+    #[structopt(name = "pid", long = "pid", parse(try_from_str = parse_hex_16))]
+    pid: Option<u16>,
+    #[structopt(name = "vid", long = "vid",  parse(try_from_str = parse_hex_16))]
+    vid: Option<u16>,
+}

src/utils.rs

@@ -0,0 +1,85 @@
+use goblin::elf::program_header::PT_LOAD;
+use rusb::GlobalContext;
+
+use std::path::PathBuf;
+use std::{fs::File, io::Read};
+
+#[derive(Debug)]
+pub enum UtilError {
+    Elf(goblin::error::Error),
+    Dfu(dfu::error::Error),
+    File(std::io::Error)
+}
+
+/// Returns a contiguous bin with 0s between non-contiguous sections and starting address from an elf.
+pub fn elf_to_bin(path: PathBuf) -> Result<(Vec<u8>, u32), UtilError> {
+    let mut file = File::open(path).map_err(|e| UtilError::File(e))?;
+    let mut buffer = vec![];
+    file.read_to_end(&mut buffer).map_err(|e| UtilError::File(e))?;
+
+    let binary = goblin::elf::Elf::parse(buffer.as_slice()).map_err(|e| UtilError::Elf(e))?;
+
+    let mut start_address: u64 = 0;
+    let mut last_address: u64 = 0;
+
+    let mut data = vec![];
+    for (i, ph) in binary
+        .program_headers
+        .iter()
+        .filter(|ph| {
+            ph.p_type == PT_LOAD
+                && ph.p_filesz > 0
+                && ph.p_offset >= binary.header.e_ehsize as u64
+                && ph.is_read()
+        })
+        .enumerate()
+    {
+        // first time through grab the starting physical address
+        if i == 0 {
+            start_address = ph.p_paddr;
+        }
+        // on subsequent passes, if there's a gap between this section and the
+        // previous one, fill it with zeros
+        else {
+            let difference = (ph.p_paddr - last_address) as usize;
+            data.resize(data.len() + difference, 0x0);
+        }
+
+        data.extend_from_slice(&buffer[ph.p_offset as usize..][..ph.p_filesz as usize]);
+
+        last_address = ph.p_paddr + ph.p_filesz;
+    }
+
+    Ok((data, start_address as u32))
+}
+
+pub fn flash_bin(
+    binary: &[u8],
+    address: u32,
+    d: &rusb::Device<GlobalContext>,
+) -> Result<(), UtilError> {
+    let mut dfu = dfu::Dfu::from_bus_device(d.bus_number(), d.address(), 0_u32, 0_u32).map_err(|e| UtilError::Dfu(e))?;
+    if binary.len() < 2048 {
+        dfu.write_flash_from_slice(address, binary).map_err(|e| UtilError::Dfu(e))?;
+    } else {
+        // hacky bug workaround
+        std::fs::write("target/out.bin", binary).map_err(|e| UtilError::File(e))?;
+        dfu.download_raw(
+            &mut std::fs::OpenOptions::new()
+                .read(true)
+                .open("target/out.bin")
+                .map_err(|e| UtilError::File(e))?,
+            address,
+            None,
+        ).map_err(|e| UtilError::Dfu(e))?;
+        std::fs::remove_file("target/out.bin").map_err(|e| UtilError::File(e))?;
+    }
+
+    Ok(())
+}
+
+pub fn vendor_map() -> std::collections::HashMap<u16, Vec<u16>> {
+    maplit::hashmap! {
+        0x0483 => vec![0xdf11],
+    }
+}