use crate::{Result, QsshError};
use std::os::unix::io::RawFd;
use tokio::io::{AsyncRead, AsyncWrite};
#[cfg(unix)]
use libc::{winsize, TIOCGWINSZ, TIOCSWINSZ};
#[derive(Debug, Clone, Copy)]
pub struct TerminalSize {
pub rows: u16,
pub cols: u16,
pub pixel_width: u16,
pub pixel_height: u16,
}
impl Default for TerminalSize {
fn default() -> Self {
Self {
rows: 24,
cols: 80,
pixel_width: 0,
pixel_height: 0,
}
}
}
pub struct Pty {
master: RawFd,
slave: RawFd,
}
impl Pty {
#[cfg(unix)]
pub fn new() -> Result<Self> {
use libc::{openpty, grantpt, unlockpt};
use std::ptr;
let mut master: RawFd = 0;
let mut slave: RawFd = 0;
unsafe {
if openpty(&mut master, &mut slave, ptr::null_mut(), ptr::null_mut(), ptr::null_mut()) != 0 {
return Err(QsshError::Io(std::io::Error::other(
"Failed to allocate PTY"
)));
}
if grantpt(master) != 0 {
libc::close(master);
libc::close(slave);
return Err(QsshError::Io(std::io::Error::other(
"Failed to grant PTY access"
)));
}
if unlockpt(master) != 0 {
libc::close(master);
libc::close(slave);
return Err(QsshError::Io(std::io::Error::other(
"Failed to unlock PTY"
)));
}
}
Ok(Self { master, slave })
}
pub fn master_fd(&self) -> RawFd {
self.master
}
pub fn slave_fd(&self) -> RawFd {
self.slave
}
#[cfg(unix)]
pub fn set_size(&self, size: TerminalSize) -> Result<()> {
let ws = winsize {
ws_row: size.rows,
ws_col: size.cols,
ws_xpixel: size.pixel_width,
ws_ypixel: size.pixel_height,
};
unsafe {
if libc::ioctl(self.master, TIOCSWINSZ, &ws) != 0 {
return Err(QsshError::Io(std::io::Error::other(
"Failed to set terminal size"
)));
}
}
Ok(())
}
#[cfg(unix)]
pub fn get_size(&self) -> Result<TerminalSize> {
let mut ws: winsize = unsafe { std::mem::zeroed() };
unsafe {
if libc::ioctl(self.master, TIOCGWINSZ, &mut ws) != 0 {
return Err(QsshError::Io(std::io::Error::other(
"Failed to get terminal size"
)));
}
}
Ok(TerminalSize {
rows: ws.ws_row,
cols: ws.ws_col,
pixel_width: ws.ws_xpixel,
pixel_height: ws.ws_ypixel,
})
}
pub fn master_async(&self) -> (PtyReader, PtyWriter) {
let master_fd = self.master;
let reader_fd = unsafe { libc::dup(master_fd) };
let writer_fd = unsafe { libc::dup(master_fd) };
(
PtyReader { fd: reader_fd },
PtyWriter { fd: writer_fd },
)
}
}
impl Drop for Pty {
fn drop(&mut self) {
unsafe {
libc::close(self.master);
libc::close(self.slave);
}
}
}
pub struct PtyReader {
fd: RawFd,
}
impl AsyncRead for PtyReader {
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<()>> {
let async_fd = match tokio::io::unix::AsyncFd::new(self.fd) {
Ok(fd) => fd,
Err(e) => return std::task::Poll::Ready(Err(e)),
};
match async_fd.poll_read_ready(cx) {
std::task::Poll::Ready(Ok(mut guard)) => {
let unfilled = buf.initialize_unfilled();
match unsafe {
libc::read(
self.fd,
unfilled.as_mut_ptr() as *mut libc::c_void,
unfilled.len(),
)
} {
-1 => {
let err = std::io::Error::last_os_error();
if err.kind() == std::io::ErrorKind::WouldBlock {
guard.clear_ready();
std::task::Poll::Pending
} else {
std::task::Poll::Ready(Err(err))
}
}
n => {
buf.advance(n as usize);
std::task::Poll::Ready(Ok(()))
}
}
}
std::task::Poll::Ready(Err(e)) => std::task::Poll::Ready(Err(e)),
std::task::Poll::Pending => std::task::Poll::Pending,
}
}
}
pub struct PtyWriter {
fd: RawFd,
}
impl AsyncWrite for PtyWriter {
fn poll_write(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
buf: &[u8],
) -> std::task::Poll<std::io::Result<usize>> {
let async_fd = match tokio::io::unix::AsyncFd::new(self.fd) {
Ok(fd) => fd,
Err(e) => return std::task::Poll::Ready(Err(e)),
};
match async_fd.poll_write_ready(cx) {
std::task::Poll::Ready(Ok(mut guard)) => {
match unsafe {
libc::write(
self.fd,
buf.as_ptr() as *const libc::c_void,
buf.len(),
)
} {
-1 => {
let err = std::io::Error::last_os_error();
if err.kind() == std::io::ErrorKind::WouldBlock {
guard.clear_ready();
std::task::Poll::Pending
} else {
std::task::Poll::Ready(Err(err))
}
}
n => std::task::Poll::Ready(Ok(n as usize)),
}
}
std::task::Poll::Ready(Err(e)) => std::task::Poll::Ready(Err(e)),
std::task::Poll::Pending => std::task::Poll::Pending,
}
}
fn poll_flush(
self: std::pin::Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<std::io::Result<()>> {
std::task::Poll::Ready(Ok(()))
}
fn poll_shutdown(
self: std::pin::Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<std::io::Result<()>> {
unsafe {
libc::close(self.fd);
}
std::task::Poll::Ready(Ok(()))
}
}
impl Drop for PtyReader {
fn drop(&mut self) {
unsafe {
libc::close(self.fd);
}
}
}
impl Drop for PtyWriter {
fn drop(&mut self) {
unsafe {
libc::close(self.fd);
}
}
}