Rust: how to build a hashmap of functions and mutate it at runtime

Question

I'm trying to parse a bunch of network packets captured in the network and perform some analysis. But I'm having a hard time dealing with procotols with dynamic port mechanism...

Some backgrounds on what "dynamic port" is: For some network protocols like FTP, 2 separate TCP channels are used for communication(command channel and data channel). Command channel uses a well-known port (like 21 for FTP), but the port that data channel uses is "dynamic": it's randomly selected by the server and wrapped in the server response for PASV command.

Typical communication procedure:

1. client establishes command channel( connects to well-known server port X)
2. client asks for the port of the data channel
3. server sends back a port Y
4. client establishes data channel (connects to server port Y)
5. client sends command on command channel(port X), server responses data on data channel(port Y)

When I'm parsing TCP packets, I use src and dst port to decide which parser function to use. To handle this situation, I think I need to keep a HashMap<u16, Parser>, mapping port to parser functions. But I just cannot make it work(mostly confusing lifetime issues). I think the main problem is: the parser functions will need to update the Hashmap that contains it. But I don't know how to solve this...

Here's what I came up with right now(still does not compile), could anyone offer some better ideas? Am I doing this the wrong way? Thanks!

use bytes::{BufMut, BytesMut};
use nom::number::complete::{be_u16, be_u32};
use nom::IResult;
use nom::bytes::complete::take;
use std::collections::HashMap;

#[derive(Debug)]
struct Command<'a> {
    command: &'a [u8],
    port: u16,
}

#[derive(Debug)]
struct Data<'a> {
    path: &'a [u8],
}

fn parse_command<'a>(input: &'a [u8]) -> IResult<&'a [u8], Command<'a>> {
    let (input, command_len) = be_u32(input)?;
    let (input, command) = take(command_len as usize)(input)?;
    let (input, port) = be_u16(input)?;
    Ok((input, Command { command, port }))
}

fn parse_data<'a>(input: &'a [u8]) -> IResult<&'a [u8], Data<'a>> {
    let (input, len) = be_u32(input)?;
    let (input, path) = take(len as usize)(input)?;
    Ok((input, Data { path }))
}

fn make_command_bytes(command: &[u8], data_port: u16) -> BytesMut {
    let mut buf = BytesMut::with_capacity(1024);
    buf.put_u32(command.len() as u32);
    buf.put(&command[..]);
    buf.put_u16(data_port);
    buf
}

fn make_data_bytes(path: &[u8]) -> BytesMut {
    let mut buf = BytesMut::with_capacity(1024);
    buf.put_u32(path.len() as u32);
    buf.put(&path[..]);
    buf
}

#[derive(Debug)]
enum Payload<'a> {
    Command(Command<'a>),
    Data(Data<'a>),
}

trait Parser<'a> {
    type Output;
    fn parse(&mut self, input: &'a [u8]) -> IResult<&'a [u8], Self::Output>;
}

struct CommandParser<'a, 'b> {
    map: &'b mut HashMap<u16, Box<dyn Parser<'a, Output = Payload<'a>>>>,
}

struct DataParser<'a, 'b> {
    map: &'b mut HashMap<u16, Box<dyn Parser<'a, Output = Payload<'a>>>>,
}

impl<'a, 'b> Parser<'a> for CommandParser<'a, 'b> {
    type Output = Payload<'a>;
    fn parse(&mut self, input: &'a [u8]) -> IResult<&'a [u8], Self::Output> {
        let (input, command) = parse_command(input)?;
        self.map
            .insert(command.port, Box::new(DataParser { map: self.map }))
            .unwrap();
        Ok((input, Payload::Command(command)))
    }
}

impl<'a, 'b> Parser<'a> for DataParser<'a, 'b> {
    type Output = Payload<'a>;
    fn parse(&mut self, input: &'a [u8]) -> IResult<&'a [u8], Self::Output> {
        let (input, data) = parse_data(input)?;
        Ok((input, Payload::Data(data)))
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let data_port = 1234;
    let command = b"test";
    let data = make_command_bytes(&command[..], data_port);

    let mut parser_map = HashMap::<u16, Box<dyn Parser<Output = Payload>>>::new();
    parser_map.insert(21, Box::new(CommandParser { map: &parser_map }));

    let parser = parser_map.get(&21).unwrap();
    match parser.parse(&data) {
        Ok((input, result)) => {
            println!("result: {:?}, remain: {:?}", result, input);
        }
        Err(e) => {
            println!("{:?}", e);
        }
    }

    Ok(())
}

the compiler says:

  --> src/main.rs:91:57
   |
91 |     parser_map.insert(21, Box::new(CommandParser { map: &parser_map }));
   |                                                         ^^^^^^^^^^^ types differ in mutability
   |
   = note: expected mutable reference `&mut HashMap<u16, Box<(dyn Parser<'_, Output = Payload<'_>> + 'static)>>`
                      found reference `&HashMap<u16, Box<dyn Parser<'_, Output = Payload<'_>>>>`

where does + 'static come from?

Answer 1

As Jmb stated in the comments, you cant store the reference in each struct. So I would keep the reference just in the main function and mutate the map from there:

match parser.parse(&data) {
    // divide `Payload::Data` and `Payload::Command`
    Ok((input, Payload::Data(data))) => {
        println!("data: {:?}, remain: {:?}", data, input);
    }
    // destructure `Command` into `command` and `port`
    Ok((input, Payload::Command(Command { command, port }))) => {
        println!("command: {:?} for {}, remain: {:?}", command, port, input);
        parser_map.insert(port, Box::new(DataParser {})).unwrap();
    }
    Err(e) => {
        println!("{:?}", e);
    }
}

Now you can remove the map from both CommandParser and DataParser and remove the unused parameters from the struct and impls:

struct CommandParser {}

struct DataParser {}

impl<'a> Parser<'a> for CommandParser {
    type Output = Payload<'a>;
    fn parse(&mut self, input: &'a [u8]) -> IResult<&'a [u8], Self::Output> {
        let (input, command) = parse_command(input)?;
        Ok((input, Payload::Command(command)))
    }
}

impl<'a> Parser<'a> for DataParser {
    type Output = Payload<'a>;
    fn parse(&mut self, input: &'a [u8]) -> IResult<&'a [u8], Self::Output> {
        let (input, data) = parse_data(input)?;
        Ok((input, Payload::Data(data)))
    }
}

Now there is only an Error about mutability:

error[E0596]: cannot borrow `**parser` as mutable, as it is behind a `&` reference
   --> src/main.rs:118:11
    |
117 |     let parser = parser_map.get(&21).unwrap();
    |         ------ help: consider changing this to be a mutable reference: `&mut Box<dyn Parser<'_, Output = Payload<'_>>>`
118 |     match parser.parse(&data) {
    |           ^^^^^^^^^^^^^^^^^^^ `parser` is a `&` reference, so the data it refers to cannot be borrowed as mutable

This can be fixed by changing let parser = parser_map.get(&21).unwrap(); to

let parser = parser_map.get_mut(&21).unwrap();

Note that the unwrap on parser_map.insert(port, Box::new(DataParser {})).unwrap(); panics because 1234 was not a port in the map before. Here is the full code

I personally would remove trait Parser and both structs Commandparser and DataParser and store just both functions in the parser_map.