Wednesday, 13 May 2015

Rust: send and receive on localhost with UDP

UDP it is the simple cousin of TCP. You send a message off into the ether, you might listen a period of time for a response, and the payloads are small enough to fit in a packet.

This post looks at sending and receiving packets in Rust.

use std::thread;
use std::net;

fn socket(listen_on: net::SocketAddr) -> net::UdpSocket {
  let attempt = net::UdpSocket::bind(listen_on);
  let mut socket;
  match attempt {
    Ok(sock) => {
      println!("Bound socket to {}", listen_on);
      socket = sock;
    },
    Err(err) => panic!("Could not bind: {}", err)
  }
  socket
}

fn read_message(socket: net::UdpSocket) -> Vec<u8> {
  let mut buf: [u8; 1] = [0; 1];
  println!("Reading data");
  let result = socket.recv_from(&mut buf);
  drop(socket);
  let mut data;
  match result {
    Ok((amt, src)) => {
      println!("Received data from {}", src);
      data = Vec::from(&buf[0..amt]);
    },
    Err(err) => panic!("Read error: {}", err)
  }
  data
}

pub fn send_message(send_addr: net::SocketAddr, target: net::SocketAddr, data: Vec<u8>) {
  let socket = socket(send_addr);
  println!("Sending data");
  let result = socket.send_to(&data, target);
  drop(socket);
  match result {
    Ok(amt) => println!("Sent {} bytes", amt),
    Err(err) => panic!("Write error: {}", err)
  }
}

pub fn listen(listen_on: net::SocketAddr) -> thread::JoinHandle<Vec<u8>> {
  let socket = socket(listen_on);
  let handle = thread::spawn(move || {
    read_message(socket)
  });
  handle
}

#[cfg(test)]
mod test {
  use std::net;
  use std::thread;
  use super::*;

  #[test]
  fn test_udp() {
    println!("UDP");
    let ip = net::Ipv4Addr::new(127, 0, 0, 1);
    let listen_addr = net::SocketAddrV4::new(ip, 8888);
    let send_addr = net::SocketAddrV4::new(ip, 8889);
    let future = listen(net::SocketAddr::V4(listen_addr));
    let message: Vec<u8> = vec![10];
 // give the thread 3s to open the socket
    thread::sleep_ms(3000);
    send_message(net::SocketAddr::V4(send_addr), net::SocketAddr::V4(listen_addr), message);
    println!("Waiting");
    let received = future.join().unwrap();
    println!("Got {} bytes", received.len());
    assert_eq!(1, received.len());
    assert_eq!(10, received[0]);
  }
}

The unit test starts listening for a UDP message in a new thread. Then it sends a UDP message to that listener. Then it consumes that message and checks it received what it sent.

The thread::sleep_ms call is used to prevent a race condition between the send and receive that will cause the unit test to fail or hang. Causing a thread to wait in this manner is inherently unreliable so avoid it in production code.

The Rust API does not currently offer a way to set UDP timeouts. Hopefully this will be added soon.

This code uses println! macros to help visualise what order things happen in. Cargo usually swallows this output but it can be observed by running tests using:

cargo test -- --nocapture

This post is just a code snippet written by someone getting started. No promises are made about code quality. Version: rustc 1.0.0-beta.4 (850151a75 2015-04-30) (built 2015-04-30)

No comments:

Post a Comment

All comments are moderated