NLinker icon

Rust select/fork example

NLinker | PRO | 01/28/21 10:57:39 AM UTC | 0 ⭐ | 8959 👁️ | Never ⏰ | []
Rust |

5.23 KB

|

None

|

0 👍

/

0 👎

extern crate libc;
 
use std::io::prelude::*;
use std::net::{TcpStream, TcpListener, UdpSocket};
use std::{io, mem, ptr, time};
use std::os::unix::io::{AsRawFd, RawFd};
use std::thread;
 
use nix::unistd::{fork, ForkResult};
 
pub struct FdSet(libc::fd_set);
 
impl FdSet {
    pub fn new() -> FdSet {
        unsafe {
            let mut raw_fd_set = mem::MaybeUninit::<libc::fd_set>::uninit();
            libc::FD_ZERO(raw_fd_set.as_mut_ptr());
            FdSet(raw_fd_set.assume_init())
        }
    }
    pub fn clear(&mut self, fd: RawFd) {
        unsafe { libc::FD_CLR(fd, &mut self.0) }
    }
    pub fn set(&mut self, fd: RawFd) {
        unsafe { libc::FD_SET(fd, &mut self.0) }
    }
    pub fn is_set(&mut self, fd: RawFd) -> bool {
        unsafe { libc::FD_ISSET(fd, &mut self.0) }
    }
}
 
fn to_fdset_ptr(opt: Option<&mut FdSet>) -> *mut libc::fd_set {
    match opt {
        None => ptr::null_mut(),
        Some(&mut FdSet(ref mut raw_fd_set)) => raw_fd_set,
    }
}
 
fn to_ptr<T>(opt: Option<&T>) -> *const T {
    match opt {
        None => ptr::null::<T>(),
        Some(p) => p,
    }
}
 
pub fn make_timeval(duration: time::Duration) -> libc::timeval {
    libc::timeval {
        tv_sec: duration.as_secs() as i64,
        tv_usec: duration.subsec_micros() as i64,
    }
}
 
pub fn select(
    nfds: libc::c_int,
    readfds: Option<&mut FdSet>,
    writefds: Option<&mut FdSet>,
    errorfds: Option<&mut FdSet>,
    timeout: Option<&libc::timeval>,
) -> io::Result<usize> {
    match unsafe {
        libc::select(
            nfds,
            to_fdset_ptr(readfds),
            to_fdset_ptr(writefds),
            to_fdset_ptr(errorfds),
            to_ptr::<libc::timeval>(timeout) as *mut libc::timeval,
        )
    } {
        -1 => Err(io::Error::last_os_error()),
        res => Ok(res as usize),
    }
}
 
pub fn main(){
    let mut msg: [u8; 128] = [0; 128];
    
    #[allow(unused_unsafe)]
    match unsafe{fork()} {
        Ok(ForkResult::Parent { child, .. }) => {
            println!("Waiting for a while on the parent process so that listener can start on child process with pid: {}!", child);
            thread::sleep(time::Duration::from_secs(10));
            println!("Trying to connect to listener in parent. \n");
            let mut tcp_stream = TcpStream::connect("127.0.0.1:8080").expect("Error connecting to listener!");
 
            let udp_socket = UdpSocket::bind("127.0.0.1:8080").expect("Failed creating a client side udp socket.");
            udp_socket.connect("127.0.0.1:8081").expect("Udp connection failure");
 
            let rfd1 = tcp_stream.as_raw_fd();
            let rfd2 = udp_socket.as_raw_fd();
 
            println!("TCP Socket: {}", rfd1);
            println!("UDP Socket: {}", rfd2);
 
            let max_fd = rfd1.max(rfd2);
            
            let mut fd_set = FdSet::new();
            fd_set.set(rfd1);
            fd_set.set(rfd2);
            match select(
                max_fd + 1,
                Some(&mut fd_set),                               // read
                None,                                            // write
                None,                                            // error
                Some(&make_timeval(time::Duration::new(10, 0))), // timeout
            ) {
                Ok(res) => {
                    println!("Select result: {}", res);
        
                    if (fd_set).is_set(rfd1) {
                        println!("TCP socket recieved something!");
                        let bytes = tcp_stream.read(&mut msg).expect("Failed to recieve message from listener!").expect("Error reading message");
                        println!("TCP Socket Recieved: {}", String::from_utf8_lossy(&msg[..bytes]));
                    }
                    if (fd_set).is_set(rfd2) {
                        println!("UDP socket recieved something!");
                        let bytes = udp_socket.recv(&mut msg).expect("Error reading message");
                        println!("UDP Socket Recieved: {}", String::from_utf8_lossy(&msg[..bytes]));
                    }
                }
                Err(err) => {
                    println!("Failed to select: {:?}", err);
                }
            }
        }
        Ok(ForkResult::Child) => {
            println!("Starting listener in child process");
 
            let tcp_listener = TcpListener::bind("127.0.0.1:8080").expect("Error starting a TCP listener!");
            println!("TCP Listener started in child process.\n");
 
            let socket = UdpSocket::bind("127.0.0.1:8081").expect("Error starting a UDP socket");
            println!("UDP socket started in child process.\n");
        
            socket.connect("127.0.0.1:8080").expect("Connect to socket failed.");
            socket.send("Hello socket!".as_bytes()).expect("data broadcast on udp failed");
 
            println!("Listening for tcp connections");
            let mut stream = tcp_listener.incoming().next().expect("Failed to find a peer").expect("Failed to connect to reciever");
            stream.write("hello world!".as_bytes()).expect("Failed to respond");
               
        }
        Err(_) => println!("Fork failed"),
    }
}

Comments

  •  icon
    01/01/70 12:00:00 AM UTC
    Plain Text |

    0 B

    |

    👍

    /

    👎

    
        
  •  icon
    01/01/70 12:00:00 AM UTC
    Plain Text |

    0 B

    |

    👍

    /

    👎