4
\$\begingroup\$

I'm new to writing rust code, and looking for any alternate approaches to writing this code or changes to make it more idiomatic, and readable.

My use of traits, specifically, I dislike. It feels like I've reached for a chainsaw to sharpen a pencil. But was the only way I could find of move closure noise outside of the main function.

But every choice is fair game to critique.

The application is a small CLI that sends a Wake-on-LAN magic packet. You can either pass the MAC address in, or you can provide an alias which gets resolved into a MAC address. So you can do something like wake macbook instead of wake 3d:4f:23:a1:c6:7a.

To map aliases to a MAC address, I've added my own field MACAddress to my ~/.ssh/config for each host. So my config looks something like this:

Host *
 IgnoreUnknown MACAddress

Host macbook
 HostName 10.0.0.100
 Port 29228
 MACAddress 3d:4f:23:a1:c6:7a
 User username

Here's the rust code:

use std::fmt::Display;

use regex::Regex;

use docopt::Docopt;
use serde::Deserialize;
use ssh2_config::{HostParams, ParseRule, SshConfig, SshParserError};
use wake_on_lan::MagicPacket;

const USAGE: &'static str = "
Wake brings together your ssh config and WakeOnLan together,
so you can use the same aliases you use to ssh to your machines,
to wake them up.

Usage:
  wake <alias>
  wake -h | --help
  wake --version

Options:
  -h --help     Show this screen.
  --version     Show version.
";

#[derive(Deserialize)]
struct Args {
    arg_alias: String,
}

const VERSION: &str = env!("CARGO_PKG_VERSION");

fn parse_mac_address(address: &str) -> Option<[u8; 6]> {
    let re = Regex::new(
        r"(?x)
        ^
        ([0-9A-f]{2})[:-]
        ([0-9A-f]{2})[:-]
        ([0-9A-f]{2})[:-]
        ([0-9A-f]{2})[:-]
        ([0-9A-f]{2})[:-]
        ([0-9A-f]{2})
        $
    ",
    )
    .unwrap();

    let (_, groups) = re.captures(address)?.extract();
    let digits = groups.map(|g| u8::from_str_radix(g, 16).unwrap());

    Some(digits)
}

fn parse_host_ssh_config(name: &str) -> Result<HostParams, SshParserError> {
    let rule = ParseRule::ALLOW_UNKNOWN_FIELDS;
    match SshConfig::parse_default_file(rule) {
        Ok(config) => return Ok(config.query(name)),
        Err(err) => return Err(err),
    }
}

trait OptionExit<T> {
    fn unwrap_or_exit(self, msg: &str) -> T;
}

impl<T> OptionExit<T> for Option<T> {
    fn unwrap_or_exit(self, msg: &str) -> T {
        match self {
            Some(var) => var,
            None => {
                eprintln!("{msg}");
                std::process::exit(1)
            }
        }
    }
}

trait ResultExit<T> {
    fn unwrap_or_exit(self) -> T;
}

impl<T, E> ResultExit<T> for Result<T, E>
where
    E: Display,
{
    fn unwrap_or_exit(self) -> T {
        self.unwrap_or_else(|err| {
            eprintln!("{}", err.to_string());
            std::process::exit(1);
        })
    }
}

fn main() {
    let docopt = Docopt::new(USAGE).unwrap();
    let args: Args = docopt
        .help(true)
        .version(Some(VERSION.to_string()))
        .deserialize()
        .unwrap_or_else(|e| e.exit());

    let alias = args.arg_alias;

    if let Some(mac_address) = parse_mac_address(&alias) {
        let _ = MagicPacket::new(&mac_address).send().unwrap_or_exit();
        println!("Magic packet sent to {alias}");
    } else {
        let host_config = parse_host_ssh_config(&alias).unwrap_or_exit();

        let field = host_config
            .ignored_fields
            .get("MACAddress")
            .unwrap_or_exit("No 'MacAddress' field defined for '{alias}' in SSH config");

        let mac_address = field
            .get(0)
            .unwrap_or_exit("No 'MacAddress' field defined for '{alias}' in SSH config");

        let mac_address = parse_mac_address(&mac_address)
            .unwrap_or_exit("'MacAddress' of '{mac_address}' for '{alias}' doesn't look like a MAC address in SSH config");

        let _ = MagicPacket::new(&mac_address).send().unwrap_or_exit();

        println!("Magic packet sent to {alias}");
    }
}

Here's the Cargo.toml:

[package]
name = "wake-cli"
version = "0.1.0"
edition = "2021"

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

[dependencies]
docopt = "1.1.1"
regex = "1.10.4"
resolve-path = "0.1.0"
serde = "1.0.197"
ssh2-config = "0.2.3"
wake-on-lan = "0.2.0"

UPDATE:

Based on comment:

use std::fmt;

use docopt::Docopt;
use lazy_regex::{lazy_regex, Lazy, Regex};
use serde::Deserialize;
use ssh2_config::{ParseRule, SshConfig};
use wake_on_lan::MagicPacket;

const USAGE: &'static str = "
Wake brings together your ssh config and WakeOnLan together,
so you can use the same aliases you use to ssh to your machines,
to wake them up.

Usage:
  wake <alias>
  wake -h | --help
  wake --version

Options:
  -h --help     Show this screen.
  --version     Show version.
";

#[derive(Deserialize)]
struct Args {
    arg_alias: String,
}

const VERSION: &str = env!("CARGO_PKG_VERSION");
const MAC_ADDRESS: Lazy<Regex> = lazy_regex!(
    r"(?x)
    ^
    ([0-9A-f]{2})[:-]
    ([0-9A-f]{2})[:-]
    ([0-9A-f]{2})[:-]
    ([0-9A-f]{2})[:-]
    ([0-9A-f]{2})[:-]
    ([0-9A-f]{2})
    $
"
);

fn parse_mac_address(address: &str) -> Option<[u8; 6]> {
    let (_, groups) = MAC_ADDRESS.captures(address)?.extract();
    let digits = groups.map(|g| u8::from_str_radix(g, 16).unwrap());

    Some(digits)
}

#[derive(Debug)]
struct SshMacParseError(String);

impl fmt::Display for SshMacParseError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.0)
    }
}

fn get_mac_from_ssh_config(alias: &str) -> Result<[u8; 6], SshMacParseError> {
    let rule = ParseRule::ALLOW_UNKNOWN_FIELDS;
    let host = SshConfig::parse_default_file(rule)
        .map(|cfg| cfg.query(alias))
        .map_err(|e| SshMacParseError(e.to_string()))?;

    let addr = host
        .ignored_fields
        .get("MACAddress")
        .ok_or(SshMacParseError("No 'MacAddress' field defined for '{alias}' in SSH config".to_string()))?
        .get(0)
        .ok_or(SshMacParseError("No 'MacAddress' value defined for '{alias}' in SSH config".to_string()))?;

    parse_mac_address(addr)
        .ok_or(SshMacParseError("'MacAddress' of '{mac_address}' for '{alias}' doesn't look like a MAC address in SSH config".to_string()))
}

fn exit<T>(msg: T) -> !
where
    T: fmt::Display,
{
    eprintln!("{}", msg.to_string());
    std::process::exit(1)
}

fn main() {
    let docopt = Docopt::new(USAGE).unwrap();
    let args: Args = docopt
        .help(true)
        .version(Some(VERSION.to_string()))
        .deserialize()
        .unwrap_or_else(|e| e.exit());

    let alias = args.arg_alias;
    let mac_address: [u8; 6];

    if let Some(_mac_address) = parse_mac_address(&alias) {
        mac_address = _mac_address;
    } else {
        mac_address = get_mac_from_ssh_config(&alias).unwrap_or_else(|e| exit(e));
    };

    MagicPacket::new(&mac_address)
        .send()
        .unwrap_or_else(|e| exit(e));

    println!("Magic packet sent to {alias}");
}
\$\endgroup\$

1 Answer 1

5
+100
\$\begingroup\$

The code looks very good to me. My suggestions are more opinion than anything else.

  1. You don't use the resolve-path crate.

  2. In your program it doesn't matter since parse_mac_address is only called once, but idiomatic usage of Regex is to have a init-once live-forever instance of the regex.

  3. parse_host_ssh_config can use Result::map instead of matching. It is also unidiomatic to put the returns in the match arms: if you omit the return, you'd get the Result as the expression result, and thus the function result.

  4. Instead of the extension traits and unwrap_or_exit thing, just have main call a run that returns a Result and uses ok_or on Options and ? for early return.

  5. Consider restructuring the key code to something like

    let mac_address = if let Some(mac_address) = parse_mac_address(&alias) { mac_address }
                      else { get_mac_from_ssh_config(&alias)? };
    let _ = MagicPacket::new(&mac_address).send()?;
    println!("Magic packet sent to {alias}");
    
\$\endgroup\$
1
  • \$\begingroup\$ Thanks, I very much appreciate the input. 2. yeah, I instinctually wanted to make it global as well, and had to find a way to do that. 3. Yes, this is much nicer and more condense. 4. I wanted to keep the 'handled' errors (so there are no 'panics!' on expected failures, and 'good' error messages). I've refactored and made much more use of ok_or, and it feels better. 5. Nice. I didn't use exactly this, as to me I think that syntax is soo rusty, that someone unfamiliar with rust would struggle to decipher it (it took me a minute). Will edit my question with new code. \$\endgroup\$
    – freebie
    Apr 18 at 22:15

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.