use std::io::{Read, Write};
use std::net::TcpStream;
use std::os::unix::net::UnixStream;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use rustls::pki_types::ServerName;
use rustls::{ClientConnection, Connection, ServerConnection};
use sha2::{Digest, Sha256};
use super::{read_frame, write_frame, CertId, PeerIdentity, RpcTransport};
use crate::rpc::{RpcError, RpcResult};
const TLS_READ_CHUNK: usize = 16 * 1024;
pub trait TlsStream: Send + Sync {
fn read(&self, buf: &mut [u8]) -> std::io::Result<usize>;
fn write(&self, buf: &[u8]) -> std::io::Result<usize>;
fn flush(&self) -> std::io::Result<()>;
fn set_read_timeout(&self, t: Option<Duration>) -> std::io::Result<()>;
fn set_write_timeout(&self, t: Option<Duration>) -> std::io::Result<()>;
}
impl TlsStream for TcpStream {
fn read(&self, buf: &mut [u8]) -> std::io::Result<usize> {
(&mut &*self).read(buf)
}
fn write(&self, buf: &[u8]) -> std::io::Result<usize> {
(&mut &*self).write(buf)
}
fn flush(&self) -> std::io::Result<()> {
(&mut &*self).flush()
}
fn set_read_timeout(&self, t: Option<Duration>) -> std::io::Result<()> {
TcpStream::set_read_timeout(self, t)
}
fn set_write_timeout(&self, t: Option<Duration>) -> std::io::Result<()> {
TcpStream::set_write_timeout(self, t)
}
}
impl TlsStream for UnixStream {
fn read(&self, buf: &mut [u8]) -> std::io::Result<usize> {
(&mut &*self).read(buf)
}
fn write(&self, buf: &[u8]) -> std::io::Result<usize> {
(&mut &*self).write(buf)
}
fn flush(&self) -> std::io::Result<()> {
(&mut &*self).flush()
}
fn set_read_timeout(&self, t: Option<Duration>) -> std::io::Result<()> {
UnixStream::set_read_timeout(self, t)
}
fn set_write_timeout(&self, t: Option<Duration>) -> std::io::Result<()> {
UnixStream::set_write_timeout(self, t)
}
}
#[cfg(all(feature = "rpc-vsock", any(target_os = "linux", target_os = "android")))]
impl TlsStream for vsock::VsockStream {
fn read(&self, buf: &mut [u8]) -> std::io::Result<usize> {
(&mut &*self).read(buf)
}
fn write(&self, buf: &[u8]) -> std::io::Result<usize> {
(&mut &*self).write(buf)
}
fn flush(&self) -> std::io::Result<()> {
(&mut &*self).flush()
}
fn set_read_timeout(&self, t: Option<Duration>) -> std::io::Result<()> {
vsock::VsockStream::set_read_timeout(self, t)
}
fn set_write_timeout(&self, t: Option<Duration>) -> std::io::Result<()> {
vsock::VsockStream::set_write_timeout(self, t)
}
}
struct IoAdapter<'a>(&'a dyn TlsStream);
impl Read for IoAdapter<'_> {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
self.0.read(buf)
}
}
impl Write for IoAdapter<'_> {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
self.0.write(buf)
}
fn flush(&mut self) -> std::io::Result<()> {
self.0.flush()
}
}
pub struct TlsTransport {
conn: Mutex<Connection>,
wlock: Mutex<()>,
stream: Box<dyn TlsStream>,
peer: PeerIdentity,
desc: String,
}
fn cert_identity(
certs: Option<&[rustls::pki_types::CertificateDer<'_>]>,
subject: &str,
) -> RpcResult<CertId> {
let leaf = certs
.and_then(|c| c.first())
.ok_or(RpcError::Protocol("TLS peer presented no certificate"))?;
let mut h = Sha256::new();
h.update(leaf.as_ref());
let mut fp = [0u8; 32];
fp.copy_from_slice(&h.finalize());
Ok(CertId::new(subject.to_string(), fp))
}
fn drive_handshake(conn: &mut Connection, stream: &dyn TlsStream) -> RpcResult<()> {
let mut io = IoAdapter(stream);
while conn.is_handshaking() {
let (_rd, _wr) = conn.complete_io(&mut io)?;
}
while conn.wants_write() {
conn.write_tls(&mut io)?;
}
Ok(())
}
impl TlsTransport {
pub fn connect_stream(
stream: Box<dyn TlsStream>,
server_name: &str,
config: Arc<rustls::ClientConfig>,
) -> RpcResult<Self> {
let name = ServerName::try_from(server_name.to_string())
.map_err(|_| RpcError::Protocol("invalid TLS server name"))?;
let cc = ClientConnection::new(config, name)
.map_err(|_| RpcError::Protocol("rustls ClientConnection::new failed"))?;
let mut conn: Connection = cc.into();
drive_handshake(&mut conn, &*stream)?;
let peer = PeerIdentity::Certificate(cert_identity(conn.peer_certificates(), server_name)?);
Ok(TlsTransport {
conn: Mutex::new(conn),
wlock: Mutex::new(()),
stream,
peer,
desc: format!("tls:{server_name}"),
})
}
pub fn accept_stream(
stream: Box<dyn TlsStream>,
config: Arc<rustls::ServerConfig>,
) -> RpcResult<Self> {
let sc = ServerConnection::new(config)
.map_err(|_| RpcError::Protocol("rustls ServerConnection::new failed"))?;
let mut conn: Connection = sc.into();
drive_handshake(&mut conn, &*stream)?;
let peer = match conn.peer_certificates() {
Some(c) if !c.is_empty() => {
PeerIdentity::Certificate(cert_identity(Some(c), "<mtls-client>")?)
}
_ => PeerIdentity::Anonymous,
};
Ok(TlsTransport {
conn: Mutex::new(conn),
wlock: Mutex::new(()),
stream,
peer,
desc: "tls:server".to_string(),
})
}
pub fn connect(
tcp: TcpStream,
server_name: &str,
config: Arc<rustls::ClientConfig>,
) -> RpcResult<Self> {
tcp.set_nodelay(true)?;
Self::connect_stream(Box::new(tcp), server_name, config)
}
pub fn accept(tcp: TcpStream, config: Arc<rustls::ServerConfig>) -> RpcResult<Self> {
tcp.set_nodelay(true)?;
Self::accept_stream(Box::new(tcp), config)
}
fn write_socket_locked(&self, cipher: &[u8]) -> RpcResult<()> {
if cipher.is_empty() {
return Ok(());
}
let mut off = 0;
while off < cipher.len() {
match self.stream.write(&cipher[off..]) {
Ok(0) => return Err(RpcError::PeerClosed),
Ok(n) => off += n,
Err(e) => return Err(e.into()),
}
}
self.stream.flush()?;
Ok(())
}
fn flush_control(&self) -> RpcResult<()> {
let Ok(_g) = self.wlock.try_lock() else {
return Ok(());
};
let cipher = {
let mut c = self.conn.lock().expect("tls conn poisoned");
if !c.wants_write() {
return Ok(());
}
let mut v = Vec::new();
c.write_tls(&mut v)?;
v
};
self.write_socket_locked(&cipher)
}
fn pump_incoming(&self) -> RpcResult<bool> {
let mut tmp = [0u8; TLS_READ_CHUNK];
let k = self.stream.read(&mut tmp)?;
if k == 0 {
return Ok(false); }
let mut c = self.conn.lock().expect("tls conn poisoned");
let mut src: &[u8] = &tmp[..k];
while !src.is_empty() {
let n = c.read_tls(&mut src)?;
if n == 0 {
break;
}
c.process_new_packets()
.map_err(|_| RpcError::Protocol("TLS record processing failed"))?;
}
Ok(true)
}
}
impl RpcTransport for TlsTransport {
fn send_frame(&self, buf: &[u8]) -> RpcResult<()> {
write_frame(&mut RawIo(self), buf)
}
fn recv_frame(&self) -> RpcResult<Vec<u8>> {
read_frame(&mut RawIo(self))
}
fn send_raw(&self, buf: &[u8]) -> RpcResult<()> {
let _g = self.wlock.lock().expect("tls wlock poisoned");
let cipher = {
let mut c = self.conn.lock().expect("tls conn poisoned");
c.writer().write_all(buf)?;
let mut v = Vec::new();
c.write_tls(&mut v)?;
v
};
self.write_socket_locked(&cipher)
}
fn recv_raw(&self, out: &mut [u8]) -> RpcResult<usize> {
loop {
{
let mut c = self.conn.lock().expect("tls conn poisoned");
match c.reader().read(out) {
Ok(n) if n > 0 => return Ok(n),
Ok(_) => return Ok(0), Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {}
Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(0),
Err(e) => return Err(e.into()),
}
}
self.flush_control()?;
if !self.pump_incoming()? {
return Ok(0);
}
}
}
fn peer_identity(&self) -> PeerIdentity {
self.peer.clone()
}
fn describe(&self) -> &str {
&self.desc
}
fn set_read_timeout(&self, timeout: Option<std::time::Duration>) -> RpcResult<()> {
self.stream.set_read_timeout(timeout)?;
Ok(())
}
fn set_write_timeout(&self, timeout: Option<std::time::Duration>) -> RpcResult<()> {
self.stream.set_write_timeout(timeout)?;
Ok(())
}
}
struct RawIo<'a>(&'a TlsTransport);
impl Read for RawIo<'_> {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
match self.0.recv_raw(buf) {
Ok(n) => Ok(n),
Err(RpcError::Io(e)) => Err(e),
Err(RpcError::PeerClosed) => Ok(0),
Err(e) => Err(std::io::Error::other(e.to_string())),
}
}
}
impl Write for RawIo<'_> {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
self.write_all(buf)?;
Ok(buf.len())
}
fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
self.0
.send_raw(buf)
.map_err(|e| std::io::Error::other(e.to_string()))
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(()) }
}