7
\$\begingroup\$

I made this simple traceroute implementation using the libpnet library, which sends ICMP echo probes with an increasing time-to-live, and then outputs the IPs of the hosts replying.

I'm new to Rust, so any style/performance tips would be appreciated!

extern crate pnet;

use std::env;
use std::error;
use std::net;
use std::str::FromStr;

use pnet::packet::icmp::echo_request::MutableEchoRequestPacket;
use pnet::packet::icmp::IcmpTypes;
use pnet::packet::ip::IpNextHeaderProtocols;
use pnet::packet::ipv4::MutableIpv4Packet;
use pnet::packet::MutablePacket;
use pnet::transport::{icmp_packet_iter, transport_channel, TransportChannelType::Layer3};
use pnet::util;

type Result<T> = std::result::Result<T, Box<error::Error>>;

static IPV4_HEADER_LEN: usize = 21;
static ICMP_HEADER_LEN: usize = 8;
static ICMP_PAYLOAD_LEN: usize = 32;

fn main() {
    std::process::exit(match run_app() {
        Ok(_) => 0,
        Err(error) => {
            eprintln!("Error: {}", error);
            1
        }
    });
}

fn run_app() -> Result<()> {
    let args: Vec<String> = env::args().collect();
    match args.len() {
        2 => {
            let protocol = Layer3(IpNextHeaderProtocols::Icmp);
            let (mut tx, mut rx) = transport_channel(1024, protocol)
                .map_err(|err| format!("Error opening the channel: {}", err))?;

            let ip_addr = net::Ipv4Addr::from_str(&args[1]).map_err(|_| "Invalid address")?;
            let mut rx = icmp_packet_iter(&mut rx);
            let mut ttl = 4;
            let mut prev_addr = None;
            loop {
                let mut buffer_ip = [0u8; 40];
                let mut buffer_icmp = [0u8; 40];
                let mut icmp_packet =
                    create_icmp_packet(&mut buffer_ip, &mut buffer_icmp, ip_addr, ttl)?;
                tx.send_to(icmp_packet, net::IpAddr::V4(ip_addr))?;
                if let Ok((_, addr)) = rx.next() {
                    if Some(addr) == prev_addr {
                        return Ok(());
                    }
                    prev_addr = Some(addr);
                    // This is not quite ideal as replies may arrive in different order 
                    // than they were sent in, but using the identification field
                    // would make this simple example more complex
                    println!("TTL: {} - {:?}", ttl, addr.to_string());
                }
                ttl += 1;
            }
        }
        _ => Err((format!("Usage: {} ip", args[0])).into()),
    }
}

fn create_icmp_packet<'a>(
    buffer_ip: &'a mut [u8],
    buffer_icmp: &'a mut [u8],
    dest: net::Ipv4Addr,
    ttl: u8,
) -> Result<MutableIpv4Packet<'a>> {
    let mut ipv4_packet = MutableIpv4Packet::new(buffer_ip).unwrap();
    ipv4_packet.set_version(4);
    ipv4_packet.set_header_length(IPV4_HEADER_LEN as u8);
    ipv4_packet.set_total_length((IPV4_HEADER_LEN + ICMP_HEADER_LEN + ICMP_PAYLOAD_LEN) as u16);
    ipv4_packet.set_ttl(ttl);
    ipv4_packet.set_next_level_protocol(IpNextHeaderProtocols::Icmp);
    ipv4_packet.set_destination(dest);

    let mut icmp_packet = MutableEchoRequestPacket::new(buffer_icmp).unwrap();
    icmp_packet.set_icmp_type(IcmpTypes::EchoRequest);
    let checksum = util::checksum(&icmp_packet.packet_mut(), 2);
    icmp_packet.set_checksum(checksum);
    ipv4_packet.set_payload(icmp_packet.packet_mut());
    Ok(ipv4_packet)
}
\$\endgroup\$
  • 1
    \$\begingroup\$ This is some good looking Rust. Have you considered reviewing other people’s questions? \$\endgroup\$ – RubberDuck Jan 6 at 1:50
1
\$\begingroup\$

I love your program’s entry point. Nearly all of my F# CLIs have an entrypoint that look nearly identical. However, Rust will automatically convert a Result into an exit code for you.

You can get rid of your run_app and wrapper by inclining it and modifying main’s signature.

fn main() -> Result<(), Box<dyn Error>> {
   // contents of `run_app` method
}

I believe this will even automatically print out the error.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

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