1
0
Fork 0
mirror of https://github.com/denoland/deno.git synced 2024-11-23 15:16:54 -05:00
denoland-deno/ext/net/raw.rs
2024-01-01 19:58:21 +00:00

303 lines
9.5 KiB
Rust

// Copyright 2018-2024 the Deno authors. All rights reserved. MIT license.
use crate::io::TcpStreamResource;
#[cfg(unix)]
use crate::io::UnixStreamResource;
use crate::ops::TcpListenerResource;
use crate::ops_tls::TlsListenerResource;
use crate::ops_tls::TlsStreamResource;
use crate::ops_tls::TLS_BUFFER_SIZE;
#[cfg(unix)]
use crate::ops_unix::UnixListenerResource;
use deno_core::error::bad_resource;
use deno_core::error::bad_resource_id;
use deno_core::error::AnyError;
use deno_core::ResourceId;
use deno_core::ResourceTable;
use deno_tls::rustls::ServerConfig;
use pin_project::pin_project;
use rustls_tokio_stream::TlsStream;
use std::rc::Rc;
use std::sync::Arc;
use tokio::net::TcpStream;
#[cfg(unix)]
use tokio::net::UnixStream;
/// A raw stream of one of the types handled by this extension.
#[pin_project(project = NetworkStreamProject)]
pub enum NetworkStream {
Tcp(#[pin] TcpStream),
Tls(#[pin] TlsStream),
#[cfg(unix)]
Unix(#[pin] UnixStream),
}
impl From<TcpStream> for NetworkStream {
fn from(value: TcpStream) -> Self {
NetworkStream::Tcp(value)
}
}
impl From<TlsStream> for NetworkStream {
fn from(value: TlsStream) -> Self {
NetworkStream::Tls(value)
}
}
#[cfg(unix)]
impl From<UnixStream> for NetworkStream {
fn from(value: UnixStream) -> Self {
NetworkStream::Unix(value)
}
}
/// A raw stream of one of the types handled by this extension.
#[derive(Copy, Clone, PartialEq, Eq)]
pub enum NetworkStreamType {
Tcp,
Tls,
#[cfg(unix)]
Unix,
}
impl NetworkStream {
pub fn local_address(&self) -> Result<NetworkStreamAddress, std::io::Error> {
match self {
Self::Tcp(tcp) => Ok(NetworkStreamAddress::Ip(tcp.local_addr()?)),
Self::Tls(tls) => Ok(NetworkStreamAddress::Ip(tls.local_addr()?)),
#[cfg(unix)]
Self::Unix(unix) => Ok(NetworkStreamAddress::Unix(unix.local_addr()?)),
}
}
pub fn peer_address(&self) -> Result<NetworkStreamAddress, std::io::Error> {
match self {
Self::Tcp(tcp) => Ok(NetworkStreamAddress::Ip(tcp.peer_addr()?)),
Self::Tls(tls) => Ok(NetworkStreamAddress::Ip(tls.peer_addr()?)),
#[cfg(unix)]
Self::Unix(unix) => Ok(NetworkStreamAddress::Unix(unix.peer_addr()?)),
}
}
pub fn stream(&self) -> NetworkStreamType {
match self {
Self::Tcp(_) => NetworkStreamType::Tcp,
Self::Tls(_) => NetworkStreamType::Tls,
#[cfg(unix)]
Self::Unix(_) => NetworkStreamType::Unix,
}
}
}
impl tokio::io::AsyncRead for NetworkStream {
fn poll_read(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
buf: &mut tokio::io::ReadBuf<'_>,
) -> std::task::Poll<std::io::Result<()>> {
match self.project() {
NetworkStreamProject::Tcp(s) => s.poll_read(cx, buf),
NetworkStreamProject::Tls(s) => s.poll_read(cx, buf),
#[cfg(unix)]
NetworkStreamProject::Unix(s) => s.poll_read(cx, buf),
}
}
}
impl tokio::io::AsyncWrite for NetworkStream {
fn poll_write(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
buf: &[u8],
) -> std::task::Poll<Result<usize, std::io::Error>> {
match self.project() {
NetworkStreamProject::Tcp(s) => s.poll_write(cx, buf),
NetworkStreamProject::Tls(s) => s.poll_write(cx, buf),
#[cfg(unix)]
NetworkStreamProject::Unix(s) => s.poll_write(cx, buf),
}
}
fn poll_flush(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), std::io::Error>> {
match self.project() {
NetworkStreamProject::Tcp(s) => s.poll_flush(cx),
NetworkStreamProject::Tls(s) => s.poll_flush(cx),
#[cfg(unix)]
NetworkStreamProject::Unix(s) => s.poll_flush(cx),
}
}
fn poll_shutdown(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), std::io::Error>> {
match self.project() {
NetworkStreamProject::Tcp(s) => s.poll_shutdown(cx),
NetworkStreamProject::Tls(s) => s.poll_shutdown(cx),
#[cfg(unix)]
NetworkStreamProject::Unix(s) => s.poll_shutdown(cx),
}
}
fn is_write_vectored(&self) -> bool {
match self {
Self::Tcp(s) => s.is_write_vectored(),
Self::Tls(s) => s.is_write_vectored(),
#[cfg(unix)]
Self::Unix(s) => s.is_write_vectored(),
}
}
fn poll_write_vectored(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
bufs: &[std::io::IoSlice<'_>],
) -> std::task::Poll<Result<usize, std::io::Error>> {
match self.project() {
NetworkStreamProject::Tcp(s) => s.poll_write_vectored(cx, bufs),
NetworkStreamProject::Tls(s) => s.poll_write_vectored(cx, bufs),
#[cfg(unix)]
NetworkStreamProject::Unix(s) => s.poll_write_vectored(cx, bufs),
}
}
}
/// A raw stream listener of one of the types handled by this extension.
pub enum NetworkStreamListener {
Tcp(tokio::net::TcpListener),
Tls(tokio::net::TcpListener, Arc<ServerConfig>),
#[cfg(unix)]
Unix(tokio::net::UnixListener),
}
pub enum NetworkStreamAddress {
Ip(std::net::SocketAddr),
#[cfg(unix)]
Unix(tokio::net::unix::SocketAddr),
}
impl NetworkStreamListener {
/// Accepts a connection on this listener.
pub async fn accept(&self) -> Result<NetworkStream, std::io::Error> {
Ok(match self {
Self::Tcp(tcp) => {
let (stream, _addr) = tcp.accept().await?;
NetworkStream::Tcp(stream)
}
Self::Tls(tcp, config) => {
let (stream, _addr) = tcp.accept().await?;
NetworkStream::Tls(TlsStream::new_server_side(
stream,
config.clone(),
TLS_BUFFER_SIZE,
))
}
#[cfg(unix)]
Self::Unix(unix) => {
let (stream, _addr) = unix.accept().await?;
NetworkStream::Unix(stream)
}
})
}
pub fn listen_address(&self) -> Result<NetworkStreamAddress, std::io::Error> {
match self {
Self::Tcp(tcp) => Ok(NetworkStreamAddress::Ip(tcp.local_addr()?)),
Self::Tls(tcp, _) => Ok(NetworkStreamAddress::Ip(tcp.local_addr()?)),
#[cfg(unix)]
Self::Unix(unix) => Ok(NetworkStreamAddress::Unix(unix.local_addr()?)),
}
}
pub fn stream(&self) -> NetworkStreamType {
match self {
Self::Tcp(..) => NetworkStreamType::Tcp,
Self::Tls(..) => NetworkStreamType::Tls,
#[cfg(unix)]
Self::Unix(..) => NetworkStreamType::Unix,
}
}
}
/// In some cases it may be more efficient to extract the resource from the resource table and use it directly (for example, an HTTP server).
/// This method will extract a stream from the resource table and return it, unwrapped.
pub fn take_network_stream_resource(
resource_table: &mut ResourceTable,
stream_rid: ResourceId,
) -> Result<NetworkStream, AnyError> {
// The stream we're attempting to unwrap may be in use somewhere else. If that's the case, we cannot proceed
// with the process of unwrapping this connection, so we just return a bad resource error.
// See also: https://github.com/denoland/deno/pull/16242
if let Ok(resource_rc) = resource_table.take::<TcpStreamResource>(stream_rid)
{
// This TCP connection might be used somewhere else.
let resource = Rc::try_unwrap(resource_rc)
.map_err(|_| bad_resource("TCP stream is currently in use"))?;
let (read_half, write_half) = resource.into_inner();
let tcp_stream = read_half.reunite(write_half)?;
return Ok(NetworkStream::Tcp(tcp_stream));
}
if let Ok(resource_rc) = resource_table.take::<TlsStreamResource>(stream_rid)
{
// This TLS connection might be used somewhere else.
let resource = Rc::try_unwrap(resource_rc)
.map_err(|_| bad_resource("TLS stream is currently in use"))?;
let (read_half, write_half) = resource.into_inner();
let tls_stream = read_half.unsplit(write_half);
return Ok(NetworkStream::Tls(tls_stream));
}
#[cfg(unix)]
if let Ok(resource_rc) = resource_table.take::<UnixStreamResource>(stream_rid)
{
// This UNIX socket might be used somewhere else.
let resource = Rc::try_unwrap(resource_rc)
.map_err(|_| bad_resource("UNIX stream is currently in use"))?;
let (read_half, write_half) = resource.into_inner();
let unix_stream = read_half.reunite(write_half)?;
return Ok(NetworkStream::Unix(unix_stream));
}
Err(bad_resource_id())
}
/// In some cases it may be more efficient to extract the resource from the resource table and use it directly (for example, an HTTP server).
/// This method will extract a stream from the resource table and return it, unwrapped.
pub fn take_network_stream_listener_resource(
resource_table: &mut ResourceTable,
listener_rid: ResourceId,
) -> Result<NetworkStreamListener, AnyError> {
if let Ok(resource_rc) =
resource_table.take::<TcpListenerResource>(listener_rid)
{
let resource = Rc::try_unwrap(resource_rc)
.map_err(|_| bad_resource("TCP socket listener is currently in use"))?;
return Ok(NetworkStreamListener::Tcp(resource.listener.into_inner()));
}
if let Ok(resource_rc) =
resource_table.take::<TlsListenerResource>(listener_rid)
{
let resource = Rc::try_unwrap(resource_rc)
.map_err(|_| bad_resource("TLS socket listener is currently in use"))?;
return Ok(NetworkStreamListener::Tls(
resource.tcp_listener.into_inner(),
resource.tls_config,
));
}
#[cfg(unix)]
if let Ok(resource_rc) =
resource_table.take::<UnixListenerResource>(listener_rid)
{
let resource = Rc::try_unwrap(resource_rc)
.map_err(|_| bad_resource("UNIX socket listener is currently in use"))?;
return Ok(NetworkStreamListener::Unix(resource.listener.into_inner()));
}
Err(bad_resource_id())
}