ptth/crates/ptth_quic/src/connection.rs

use crate::prelude::*;

pub struct NewConnection {
	pub local_send: tokio::net::tcp::OwnedWriteHalf,
	pub local_recv: tokio::net::tcp::OwnedReadHalf,
	pub relay_send: quinn::SendStream,
	pub relay_recv: quinn::RecvStream,
}

pub struct Connection {
	// Blue and green because they're not necessarily uplink nor downlink.
	// It depends on whether the client or server is using us.
	
	task_blue: JoinHandle <anyhow::Result <()>>,
	task_green: JoinHandle <anyhow::Result <()>>,
}

impl NewConnection {
	pub fn build (self) -> Connection {
		let Self {
			mut local_send,
			mut local_recv,
			mut relay_send,
			mut relay_recv,
		} = self;
		
		let task_blue = tokio::spawn (async move {
			// Downlink - Relay server to local client
			
			let mut buf = vec! [0u8; 65_536];
			while let Some (bytes_read) = relay_recv.read (&mut buf).await? {
				let buf_slice = &buf [0..bytes_read];
				trace! ("Uplink relaying {} bytes", bytes_read);
				local_send.write_all (buf_slice).await?;
			}
			
			trace! ("Downlink closed");
			
			Ok::<_, anyhow::Error> (())
		});
		
		let task_green = tokio::spawn (async move {
			// Uplink - local client to relay server
			
			let mut buf = vec! [0u8; 65_536];
			loop {
				let bytes_read = local_recv.read (&mut buf).await?;
				if bytes_read == 0 {
					break;
				}
				let buf_slice = &buf [0..bytes_read];
				trace! ("Downlink relaying {} bytes", bytes_read);
				relay_send.write_all (buf_slice).await?;
			}
			
			trace! ("Uplink closed");
			
			Ok::<_, anyhow::Error> (())
		});
		
		Connection {
			task_blue,
			task_green,
		}
	}
}

impl Connection {
	pub async fn wait_for_close (self) -> anyhow::Result <()> {
		self.task_blue.await??;
		self.task_green.await??;
		Ok (())
	}
}