diff --git a/src/clipboard.rs b/src/clipboard.rs index 5339570..5ca2335 100644 --- a/src/clipboard.rs +++ b/src/clipboard.rs @@ -1,24 +1,33 @@ use std::{ + borrow::Cow, collections::hash_map::DefaultHasher, error::Error, fmt, hash::{Hash, Hasher}, + mem, sync::atomic::{AtomicU64, Ordering}, time::Duration, }; -use arboard::Error as ClipboardError; -use tokio::{sync::Mutex, time::sleep}; +use arboard::ImageData; +use tokio::{ + io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}, + sync::Mutex, + time::sleep, +}; +use tracing::trace; pub struct Clipboard { clipboard: Mutex, - current: AtomicU64, + current_text: AtomicU64, + current_image: AtomicU64, } impl fmt::Debug for Clipboard { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("Clipboard") - .field("current", &self.current) + .field("current_text", &self.current_text) + .field("current_image", &self.current_image) .finish() } } @@ -26,41 +35,197 @@ impl fmt::Debug for Clipboard { impl Clipboard { pub fn new() -> Self { let mut clipboard = arboard::Clipboard::new().unwrap(); - let current = AtomicU64::new(hash(&clipboard.get_text().unwrap_or_default())); + let current_text = AtomicU64::new(clipboard.get_text().map(hash).unwrap_or_default()); + let current_image = AtomicU64::new( + clipboard + .get_image() + .map(|img| hash(img.bytes)) + .unwrap_or_default(), + ); Self { clipboard: Mutex::new(clipboard), - current, + current_text, + current_image, } } - pub async fn copy(&self, text: &str) -> Result<(), Box> { - let hashed = hash(text); - if self.current.load(Ordering::SeqCst) != hashed { - while let Err(ClipboardError::ClipboardOccupied) = - self.clipboard.lock().await.set_text(text) - { - sleep(Duration::from_secs(1)).await; + pub async fn copy( + &self, + obj: impl Into, + ) -> Result<(), Box> { + let obj = obj.into(); + let hashed = hash(&obj); + + match obj { + ClipboardObject::Text(text) => { + if self.current_text.load(Ordering::SeqCst) != hashed { + self.clipboard.lock().await.set_text(text)?; + self.current_text.store(hashed, Ordering::SeqCst); + } } - self.current.store(hashed, Ordering::SeqCst); - } + ClipboardObject::Image(img) => { + if self.current_image.load(Ordering::SeqCst) != hashed { + self.clipboard.lock().await.set_image(img)?; + self.current_image.store(hashed, Ordering::SeqCst); + } + } + }; Ok(()) } - pub async fn paste(&self) -> Result> { + pub async fn paste(&self) -> Result> { loop { - let paste = self.clipboard.lock().await.get_text().unwrap_or_default(); + let mut clip = self.clipboard.lock().await; + + let paste = clip.get_text().unwrap_or_default(); let hashed = hash(&paste); - if !paste.is_empty() && hashed != self.current.load(Ordering::SeqCst) { - self.current.store(hashed, Ordering::SeqCst); - break Ok(paste); + if !paste.is_empty() && hashed != self.current_text.load(Ordering::SeqCst) { + self.current_text.store(hashed, Ordering::SeqCst); + break Ok(ClipboardObject::Text(paste)); + } + + if let Ok(paste) = clip.get_image() { + let hashed = hash(&paste.bytes); + if !paste.bytes.is_empty() && hashed != self.current_image.load(Ordering::SeqCst) { + self.current_image.store(hashed, Ordering::SeqCst); + break Ok(ClipboardObject::Image(paste)); + } } sleep(Duration::from_secs(1)).await; } } } -fn hash(val: &str) -> u64 { +#[derive(Debug)] +pub enum ClipboardObject { + Text(String), + Image(ImageData<'static>), +} + +impl AsRef<[u8]> for ClipboardObject { + fn as_ref(&self) -> &[u8] { + match self { + Self::Text(txt) => txt.as_ref(), + Self::Image(img) => img.bytes.as_ref(), + } + } +} + +#[repr(u8)] +enum ClipboardObjectType { + Text = 1, + Image = 2, +} + +impl ClipboardObject { + pub async fn from_reader( + mut reader: impl AsyncRead + Send + Unpin, + ) -> Result> { + let mut buf = [0; 1]; + reader.read_exact(&mut buf).await?; + + trace!("Read kind {buf:?}"); + let kind = match buf[0] { + 1 => ClipboardObjectType::Text, + 2 => ClipboardObjectType::Image, + n => return Err(format!("Invalid clipboard object type {n}").into()), + }; + + match kind { + ClipboardObjectType::Text => { + let mut buf = [0; mem::size_of::()]; + reader.read_exact(&mut buf).await?; + let len = u64::from_be_bytes(buf).try_into()?; + trace!(len, "Read text len"); + + let mut buf = vec![0; len]; + reader.read_exact(&mut buf).await?; + trace!(len, "Read text"); + + let text = std::str::from_utf8(&buf)?; + Ok(Self::Text(text.to_string())) + } + + ClipboardObjectType::Image => { + let mut buf = [0; mem::size_of::()]; + reader.read_exact(&mut buf).await?; + let width = u64::from_be_bytes(buf).try_into()?; + trace!(width, "Read image width"); + + let mut buf = [0; mem::size_of::()]; + reader.read_exact(&mut buf).await?; + let height = u64::from_be_bytes(buf).try_into()?; + trace!(height, "Read image height"); + + let mut buf = [0; mem::size_of::()]; + reader.read_exact(&mut buf).await?; + let len = u64::from_be_bytes(buf).try_into()?; + trace!(width, height, len, "Read image metadata"); + + let mut buf = vec![0; len]; + reader.read_exact(&mut buf).await?; + trace!(width, height, len, "Read image"); + + let img = ImageData { + width, + height, + bytes: Cow::from(buf), + }; + + Ok(Self::Image(img)) + } + } + } + + pub async fn write( + self, + mut writer: impl AsyncWrite + Send + Unpin, + ) -> Result<(), Box> { + let buf = match self { + Self::Text(ref text) => { + trace!(len = text.as_bytes().len(), "Sending text"); + + [ + &[ClipboardObjectType::Text as u8][..], + &u64::try_from(text.as_bytes().len())?.to_be_bytes()[..], + ] + .concat() + } + + Self::Image(ref img) => { + trace!( + width = img.width, + height = img.height, + len = img.bytes.len(), + "Sending image" + ); + + [ + &[ClipboardObjectType::Image as u8][..], + &u64::try_from(img.width)?.to_be_bytes()[..], + &u64::try_from(img.height)?.to_be_bytes()[..], + &u64::try_from(img.bytes.len())?.to_be_bytes()[..], + ] + .concat() + } + }; + + writer.write_all(&buf).await?; + + let buf = match self { + Self::Text(ref text) => text.as_bytes(), + Self::Image(ref img) => &img.bytes, + }; + + writer.write_all(buf).await?; + trace!(len = buf.len(), "Clipboard sent"); + + Ok(()) + } +} + +fn hash(val: impl AsRef<[u8]>) -> u64 { let mut hasher = DefaultHasher::new(); - val.hash(&mut hasher); + val.as_ref().hash(&mut hasher); hasher.finish() } diff --git a/src/main.rs b/src/main.rs index 2680b88..bc8202b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,9 +1,10 @@ -use std::{error::Error, io, mem, process::exit, sync::Arc, time::Duration}; +use std::{error::Error, io, process::exit, sync::Arc, time::Duration}; use clap::{command, Parser}; +use clipboard::ClipboardObject; use rustls::{client::ServerCertVerifier, Certificate, PrivateKey, ServerName}; use tokio::{ - io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}, + io::{AsyncRead, AsyncWrite, AsyncWriteExt}, net::{TcpListener, TcpStream, UdpSocket}, select, time::{sleep, timeout}, @@ -162,14 +163,14 @@ async fn send_clipboard( mut stream: impl AsyncWrite + Send + Unpin, ) -> Result<(), Box> { loop { - let paste = clipboard.paste().in_current_span().await?; - let text = paste.as_bytes(); - let buf = [&text.len().to_be_bytes(), text].concat(); - trace!(text = paste, "Sent text"); - if stream.write(&buf).await? == 0 { - trace!("Stream closed"); - break Ok(()); - } + clipboard + .paste() + .in_current_span() + .await? + .write(&mut stream) + .in_current_span() + .await?; + stream.flush().await?; } } @@ -179,19 +180,10 @@ async fn recv_clipboard( mut stream: impl AsyncRead + Send + Unpin, ) -> Result<(), Box> { loop { - let mut buf = [0; mem::size_of::()]; - if stream.read(&mut buf).await? == 0 { - trace!("Stream closed"); - break Ok(()); - } - let len = usize::from_be_bytes(buf); - let mut buf = vec![0; len]; - stream.read_exact(&mut buf).await?; - - if let Ok(text) = std::str::from_utf8(&buf) { - trace!(text = text, "Received text"); - clipboard.copy(text).in_current_span().await?; - } + let obj = ClipboardObject::from_reader(&mut stream) + .in_current_span() + .await?; + clipboard.copy(obj).in_current_span().await?; } }