NLinker

Rust select/fork example

Jan 28th, 2021
636
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. extern crate libc;
  2.  
  3. use std::io::prelude::*;
  4. use std::net::{TcpStream, TcpListener, UdpSocket};
  5. use std::{io, mem, ptr, time};
  6. use std::os::unix::io::{AsRawFd, RawFd};
  7. use std::thread;
  8.  
  9. use nix::unistd::{fork, ForkResult};
  10.  
  11. pub struct FdSet(libc::fd_set);
  12.  
  13. impl FdSet {
  14.     pub fn new() -> FdSet {
  15.         unsafe {
  16.             let mut raw_fd_set = mem::MaybeUninit::<libc::fd_set>::uninit();
  17.             libc::FD_ZERO(raw_fd_set.as_mut_ptr());
  18.             FdSet(raw_fd_set.assume_init())
  19.         }
  20.     }
  21.     pub fn clear(&mut self, fd: RawFd) {
  22.         unsafe { libc::FD_CLR(fd, &mut self.0) }
  23.     }
  24.     pub fn set(&mut self, fd: RawFd) {
  25.         unsafe { libc::FD_SET(fd, &mut self.0) }
  26.     }
  27.     pub fn is_set(&mut self, fd: RawFd) -> bool {
  28.         unsafe { libc::FD_ISSET(fd, &mut self.0) }
  29.     }
  30. }
  31.  
  32. fn to_fdset_ptr(opt: Option<&mut FdSet>) -> *mut libc::fd_set {
  33.     match opt {
  34.         None => ptr::null_mut(),
  35.         Some(&mut FdSet(ref mut raw_fd_set)) => raw_fd_set,
  36.     }
  37. }
  38.  
  39. fn to_ptr<T>(opt: Option<&T>) -> *const T {
  40.     match opt {
  41.         None => ptr::null::<T>(),
  42.         Some(p) => p,
  43.     }
  44. }
  45.  
  46. pub fn make_timeval(duration: time::Duration) -> libc::timeval {
  47.     libc::timeval {
  48.         tv_sec: duration.as_secs() as i64,
  49.         tv_usec: duration.subsec_micros() as i64,
  50.     }
  51. }
  52.  
  53. pub fn select(
  54.     nfds: libc::c_int,
  55.     readfds: Option<&mut FdSet>,
  56.     writefds: Option<&mut FdSet>,
  57.     errorfds: Option<&mut FdSet>,
  58.     timeout: Option<&libc::timeval>,
  59. ) -> io::Result<usize> {
  60.     match unsafe {
  61.         libc::select(
  62.             nfds,
  63.             to_fdset_ptr(readfds),
  64.             to_fdset_ptr(writefds),
  65.             to_fdset_ptr(errorfds),
  66.             to_ptr::<libc::timeval>(timeout) as *mut libc::timeval,
  67.         )
  68.     } {
  69.         -1 => Err(io::Error::last_os_error()),
  70.         res => Ok(res as usize),
  71.     }
  72. }
  73.  
  74. pub fn main(){
  75.     let mut msg: [u8; 128] = [0; 128];
  76.    
  77.     #[allow(unused_unsafe)]
  78.     match unsafe{fork()} {
  79.         Ok(ForkResult::Parent { child, .. }) => {
  80.             println!("Waiting for a while on the parent process so that listener can start on child process with pid: {}!", child);
  81.             thread::sleep(time::Duration::from_secs(10));
  82.             println!("Trying to connect to listener in parent. \n");
  83.             let mut tcp_stream = TcpStream::connect("127.0.0.1:8080").expect("Error connecting to listener!");
  84.  
  85.             let udp_socket = UdpSocket::bind("127.0.0.1:8080").expect("Failed creating a client side udp socket.");
  86.             udp_socket.connect("127.0.0.1:8081").expect("Udp connection failure");
  87.  
  88.             let rfd1 = tcp_stream.as_raw_fd();
  89.             let rfd2 = udp_socket.as_raw_fd();
  90.  
  91.             println!("TCP Socket: {}", rfd1);
  92.             println!("UDP Socket: {}", rfd2);
  93.  
  94.             let max_fd = rfd1.max(rfd2);
  95.            
  96.             let mut fd_set = FdSet::new();
  97.             fd_set.set(rfd1);
  98.             fd_set.set(rfd2);
  99.             match select(
  100.                 max_fd + 1,
  101.                 Some(&mut fd_set),                               // read
  102.                 None,                                            // write
  103.                 None,                                            // error
  104.                 Some(&make_timeval(time::Duration::new(10, 0))), // timeout
  105.             ) {
  106.                 Ok(res) => {
  107.                     println!("Select result: {}", res);
  108.        
  109.                     if (fd_set).is_set(rfd1) {
  110.                         println!("TCP socket recieved something!");
  111.                         let bytes = tcp_stream.read(&mut msg).expect("Failed to recieve message from listener!").expect("Error reading message");
  112.                         println!("TCP Socket Recieved: {}", String::from_utf8_lossy(&msg[..bytes]));
  113.                     }
  114.                     if (fd_set).is_set(rfd2) {
  115.                         println!("UDP socket recieved something!");
  116.                         let bytes = udp_socket.recv(&mut msg).expect("Error reading message");
  117.                         println!("UDP Socket Recieved: {}", String::from_utf8_lossy(&msg[..bytes]));
  118.                     }
  119.                 }
  120.                 Err(err) => {
  121.                     println!("Failed to select: {:?}", err);
  122.                 }
  123.             }
  124.         }
  125.         Ok(ForkResult::Child) => {
  126.             println!("Starting listener in child process");
  127.  
  128.             let tcp_listener = TcpListener::bind("127.0.0.1:8080").expect("Error starting a TCP listener!");
  129.             println!("TCP Listener started in child process.\n");
  130.  
  131.             let socket = UdpSocket::bind("127.0.0.1:8081").expect("Error starting a UDP socket");
  132.             println!("UDP socket started in child process.\n");
  133.        
  134.             socket.connect("127.0.0.1:8080").expect("Connect to socket failed.");
  135.             socket.send("Hello socket!".as_bytes()).expect("data broadcast on udp failed");
  136.  
  137.             println!("Listening for tcp connections");
  138.             let mut stream = tcp_listener.incoming().next().expect("Failed to find a peer").expect("Failed to connect to reciever");
  139.             stream.write("hello world!".as_bytes()).expect("Failed to respond");
  140.                
  141.         }
  142.         Err(_) => println!("Fork failed"),
  143.     }
  144. }
RAW Paste Data