rustocat/src/udp.rs

use crate::listener::Listener;
use crate::shutdown::Shutdown;
use rand::seq::SliceRandom;
use std::collections::HashMap;
use std::error::Error;
use std::sync::Arc;
use tokio::net::UdpSocket;
use tokio::sync::broadcast;
use tokio::sync::mpsc::{channel, Sender};
use tokio::sync::Mutex;
use tokio::time::Instant;

async fn create_dual_udpsocket(
    bind_addr: &String,
) -> Result<(UdpSocket, UdpSocket), Box<dyn Error>> {
    let listener_std = std::net::UdpSocket::bind(bind_addr)?;
    let responder_std = listener_std.try_clone()?;
    let listener = UdpSocket::from_std(listener_std)?;
    let responder = UdpSocket::from_std(responder_std)?;

    return Ok((listener, responder));
}

fn get_udp_background_send(socket: UdpSocket, mut exit: Shutdown) -> Sender<(Vec<u8>, String)> {
    let (tx, mut rx) = channel::<(Vec<u8>, String)>(1);

    tokio::spawn(async move {
        loop {
            let (buf, dest) = (tokio::select! {
                res = rx.recv() => Some(res.unwrap()),
                _ = exit.recv() => {
                  println!("Exiting listener!");
                  return;
               }
            })
            .unwrap();

            let to_send = buf.as_slice();
            socket.send_to(to_send, &dest).await.expect(&format!(
                "Failed to forward response from upstream server to client {}",
                dest
            ));
        }
    });

    return tx;
}

struct UdpHandler {
    last_packet: Arc<Mutex<Instant>>,
    kill: broadcast::Sender<()>,
    target: String,
    sender: Sender<(Vec<u8>, String)>,
}

impl UdpHandler {
    async fn start(
        target: String,
        source: String,
        sender: Sender<(Vec<u8>, String)>,
    ) -> Result<UdpHandler, Box<dyn Error>> {
        // Kill Channel
        let (tx, mut rx) = broadcast::channel::<()>(1);

        let (listener, responder) = create_dual_udpsocket(&"0.0.0.0:0".to_owned()).await?;

        let s = get_udp_background_send(responder, Shutdown::new(tx.subscribe()));

        let last_packet = Arc::new(Mutex::new(Instant::now()));

        let handler = UdpHandler {
            kill: tx,
            last_packet: last_packet.clone(),
            // source: source.clone(),
            target: target.clone(),
            sender: s,
        };

        listener.connect(target).await?;
        tokio::spawn(async move {
            let mut buf = [0; 64 * 1024];
            loop {
                let (num_bytes, _) = tokio::select! {
                    res = listener.recv_from(&mut buf) => res.unwrap(),
                    _ = rx.recv() => {
                        // FIXME: Source of memory leaks?
                        return ;
                    }
                };
                let mut n = last_packet.lock().await;
                *n = Instant::now();
                sender
                    .send((buf[0..num_bytes].to_vec(), source.clone()))
                    .await
                    .expect(&format!("Failed to send answer to sender {}", source));
            }
        });
        return Ok(handler);
    }

    async fn exit(&self) -> Result<(), Box<dyn Error>> {
        self.kill.send(())?;

        Ok(())
    }

    async fn on_packet(&mut self, pkg: Vec<u8>) -> Result<(), Box<dyn Error>> {
        let mut n = self.last_packet.lock().await;
        *n = Instant::now();
        self.sender.send((pkg, self.target.clone())).await?;
        return Ok(());
    }
}

pub(crate) async fn start_udp_listener(
    mut listener_config: Listener,
) -> Result<(), Box<dyn Error>> {
    println!("start listening on {}", &listener_config.source);
    let (listener, responder) =
        create_dual_udpsocket(&listener_config.source)
            .await
            .expect(&format!(
                "Failed to clone primary listening address socket {}",
                &listener_config.source,
            ));

    let sender = get_udp_background_send(responder, listener_config.shutdown.clone());

    let mut connections: HashMap<String, UdpHandler> = HashMap::new();

    let mut buf = [0; 64 * 1024];
    loop {
        let (num_bytes, src_addr) = tokio::select! {
           res = listener.recv_from(&mut buf) => res?,
           _ = listener_config.shutdown.recv() => {
              println!("Exiting listener!");
              break;
           }
        };

        let addr = src_addr.to_string();
        let handler_opt = connections.get_mut(&addr);
        let vec = buf[0..num_bytes].to_vec();
        let handler = match handler_opt {
            Some(handler) => handler,
            None => {
                let targets = listener_config.targets.read().await;
                let selected_target = {
                    let mut rng = rand::thread_rng();
                    targets.choose(&mut rng).unwrap()
                };

                let handler =
                    UdpHandler::start(selected_target.clone(), addr.clone(), sender.clone())
                        .await?;
                connections.insert(addr.clone(), handler);

                connections.get_mut(&addr).unwrap()
            }
        };

        match handler.on_packet(vec).await {
            Ok(_) => (),
            Err(err) => {
                println!("Failed to forward request from client to server {}", err);
                return Ok(());
            }
        }
    }

    for handler in connections.values() {
        handler.exit().await?;
    }

    Ok(())
}