ptth/crates/udp_over_tcp/src/server.rs

use std::{
	net::{
		Ipv4Addr,
		SocketAddrV4,
	},
	sync::Arc,
};

use tokio::{
	net::{
		TcpListener,
		TcpStream,
		UdpSocket,
	},
	spawn,
};

use crate::loops;

#[derive (Clone)]
pub struct Config {
	/// The well-known TCP port that the UDP-over-TCP server will bind
	pub tcp_port: u16,
	
	/// The well-known UDP port that the PTTH_QUIC relay will bind
	pub udp_port: u16,
}

pub struct Listener {
	cfg: Config,
	tcp_listener: TcpListener,
}

impl Listener {
	pub async fn new (cfg: Config) -> anyhow::Result <Self> {
		let tcp_listener = TcpListener::bind ((Ipv4Addr::UNSPECIFIED, cfg.tcp_port)).await?;
		
		Ok (Self {
			cfg,
			tcp_listener,
		})
	}
	
	pub fn tcp_port (&self) -> anyhow::Result <u16> {
		Ok (self.tcp_listener.local_addr ()?.port ())
	}
	
	pub async fn run (self) -> anyhow::Result <()> {
		let Self {
			cfg,
			tcp_listener,
		} = self;
		
		loop { 
			let (conn, _peer_addr) = tcp_listener.accept ().await?;
			
			let cfg = cfg.clone ();
			spawn (handle_connection (cfg, conn));
		}
	}
}

async fn handle_connection (cfg: Config, conn: TcpStream) -> anyhow::Result <()> {
	let udp_sock = UdpSocket::bind (SocketAddrV4::new (Ipv4Addr::UNSPECIFIED, 0)).await?;
	udp_sock.connect ((Ipv4Addr::LOCALHOST, cfg.udp_port)).await?;
	
	let (tcp_read, tcp_write) = conn.into_split ();
	
	let rx_task;
	let tx_task;
	
	{
		let udp_sock = Arc::new (udp_sock);
		rx_task = spawn (loops::rx (Arc::clone (&udp_sock), tcp_read));
		tx_task = spawn (loops::tx (Arc::clone (&udp_sock), tcp_write));
	}
	
	tokio::select! {
		_val = tx_task => {
			println! ("server_handle_connection: tx_task exited, exiting");
		}
		_val = rx_task => {
			println! ("server_handle_connection: rx_task exited, exiting");
		}
	}
	
	Ok (())
}