tenant.openagents/omega

1use std::{
2    fs::File,
3    io::{ErrorKind, Write},
4    os::fd::{AsRawFd, BorrowedFd, OwnedFd},
5};
6
7use calloop::{LoopHandle, PostAction};
8use filedescriptor::Pipe;
9use strum::IntoEnumIterator;
10use wayland_client::{Connection, protocol::wl_data_offer::WlDataOffer};
11use wayland_protocols::wp::primary_selection::zv1::client::zwp_primary_selection_offer_v1::ZwpPrimarySelectionOfferV1;
12
13use crate::linux::{
14    WaylandClientStatePtr,
15    platform::{PIPE_READ_TIMEOUT, read_fd_with_timeout},
16};
17use gpui::{ClipboardEntry, ClipboardItem, Image, ImageFormat, hash};
18
19/// Text mime types that we'll offer to other programs.
20pub(crate) const TEXT_MIME_TYPES: [&str; 3] =
21    ["text/plain;charset=utf-8", "UTF8_STRING", "text/plain"];
22pub(crate) const FILE_LIST_MIME_TYPE: &str = "text/uri-list";
23
24/// Text mime types that we'll accept from other programs.
25pub(crate) const ALLOWED_TEXT_MIME_TYPES: [&str; 2] = ["text/plain;charset=utf-8", "UTF8_STRING"];
26
27pub(crate) struct Clipboard {
28    connection: Connection,
29    loop_handle: LoopHandle<'static, WaylandClientStatePtr>,
30    self_mime: String,
31
32    // Internal clipboard
33    contents: Option<ClipboardItem>,
34    primary_contents: Option<ClipboardItem>,
35
36    // External clipboard
37    cached_read: Option<ClipboardItem>,
38    current_offer: Option<DataOffer<WlDataOffer>>,
39    cached_primary_read: Option<ClipboardItem>,
40    current_primary_offer: Option<DataOffer<ZwpPrimarySelectionOfferV1>>,
41}
42
43pub(crate) trait ReceiveData {
44    fn receive_data(&self, mime_type: String, fd: BorrowedFd<'_>);
45}
46
47impl ReceiveData for WlDataOffer {
48    fn receive_data(&self, mime_type: String, fd: BorrowedFd<'_>) {
49        self.receive(mime_type, fd);
50    }
51}
52
53impl ReceiveData for ZwpPrimarySelectionOfferV1 {
54    fn receive_data(&self, mime_type: String, fd: BorrowedFd<'_>) {
55        self.receive(mime_type, fd);
56    }
57}
58
59#[derive(Clone, Debug)]
60/// Wrapper for `WlDataOffer` and `ZwpPrimarySelectionOfferV1`, used to help track mime types.
61pub(crate) struct DataOffer<T: ReceiveData> {
62    pub inner: T,
63    mime_types: Vec<String>,
64}
65
66impl<T: ReceiveData> DataOffer<T> {
67    pub fn new(offer: T) -> Self {
68        Self {
69            inner: offer,
70            mime_types: Vec::new(),
71        }
72    }
73
74    pub fn add_mime_type(&mut self, mime_type: String) {
75        self.mime_types.push(mime_type)
76    }
77
78    fn has_mime_type(&self, mime_type: &str) -> bool {
79        self.mime_types.iter().any(|t| t == mime_type)
80    }
81
82    fn read_bytes(&self, connection: &Connection, mime_type: &str) -> Option<Vec<u8>> {
83        let pipe = Pipe::new().unwrap();
84        self.inner.receive_data(mime_type.to_string(), unsafe {
85            BorrowedFd::borrow_raw(pipe.write.as_raw_fd())
86        });
87        let fd = pipe.read;
88        drop(pipe.write);
89
90        connection.flush().unwrap();
91
92        match read_fd_with_timeout(fd, PIPE_READ_TIMEOUT) {
93            Ok(bytes) => Some(bytes),
94            Err(err) => {
95                log::error!("error reading clipboard pipe: {err:?}");
96                None
97            }
98        }
99    }
100
101    fn read_text(&self, connection: &Connection) -> Option<ClipboardItem> {
102        let mime_type = self.mime_types.iter().find(|&mime_type| {
103            ALLOWED_TEXT_MIME_TYPES
104                .iter()
105                .any(|&allowed| allowed == mime_type)
106        })?;
107        let bytes = self.read_bytes(connection, mime_type)?;
108        let text_content = match String::from_utf8(bytes) {
109            Ok(content) => content,
110            Err(e) => {
111                log::error!("Failed to convert clipboard content to UTF-8: {}", e);
112                return None;
113            }
114        };
115
116        // Normalize the text to unix line endings, otherwise
117        // copying from eg: firefox inserts a lot of blank
118        // lines, and that is super annoying.
119        let result = text_content.replace("\r\n", "\n");
120        Some(ClipboardItem::new_string(result))
121    }
122
123    fn read_image(&self, connection: &Connection) -> Option<ClipboardItem> {
124        for format in ImageFormat::iter() {
125            let mime_type = format.mime_type();
126            if !self.has_mime_type(mime_type) {
127                continue;
128            }
129
130            if let Some(bytes) = self.read_bytes(connection, mime_type) {
131                let id = hash(&bytes);
132                return Some(ClipboardItem {
133                    entries: vec![ClipboardEntry::Image(Image { format, bytes, id })],
134                });
135            }
136        }
137        None
138    }
139}
140
141impl Clipboard {
142    pub fn new(
143        connection: Connection,
144        loop_handle: LoopHandle<'static, WaylandClientStatePtr>,
145    ) -> Self {
146        Self {
147            connection,
148            loop_handle,
149            self_mime: format!("pid/{}", std::process::id()),
150
151            contents: None,
152            primary_contents: None,
153
154            cached_read: None,
155            current_offer: None,
156            cached_primary_read: None,
157            current_primary_offer: None,
158        }
159    }
160
161    pub fn set(&mut self, item: ClipboardItem) {
162        self.contents = Some(item);
163    }
164
165    pub fn set_primary(&mut self, item: ClipboardItem) {
166        self.primary_contents = Some(item);
167    }
168
169    pub fn set_offer(&mut self, data_offer: Option<DataOffer<WlDataOffer>>) {
170        self.cached_read = None;
171        self.current_offer = data_offer;
172    }
173
174    pub fn set_primary_offer(&mut self, data_offer: Option<DataOffer<ZwpPrimarySelectionOfferV1>>) {
175        self.cached_primary_read = None;
176        self.current_primary_offer = data_offer;
177    }
178
179    pub fn self_mime(&self) -> String {
180        self.self_mime.clone()
181    }
182
183    pub fn send(&self, _mime_type: String, fd: OwnedFd) {
184        if let Some(text) = self.contents.as_ref().and_then(|contents| contents.text()) {
185            self.send_internal(fd, text.as_bytes().to_owned());
186        }
187    }
188
189    pub fn send_primary(&self, _mime_type: String, fd: OwnedFd) {
190        if let Some(text) = self
191            .primary_contents
192            .as_ref()
193            .and_then(|contents| contents.text())
194        {
195            self.send_internal(fd, text.as_bytes().to_owned());
196        }
197    }
198
199    pub fn read(&mut self) -> Option<ClipboardItem> {
200        let offer = self.current_offer.as_ref()?;
201        if let Some(cached) = self.cached_read.clone() {
202            return Some(cached);
203        }
204
205        if offer.has_mime_type(&self.self_mime) {
206            return self.contents.clone();
207        }
208
209        let item = offer
210            .read_text(&self.connection)
211            .or_else(|| offer.read_image(&self.connection))?;
212
213        self.cached_read = Some(item.clone());
214        Some(item)
215    }
216
217    pub fn read_primary(&mut self) -> Option<ClipboardItem> {
218        let offer = self.current_primary_offer.as_ref()?;
219        if let Some(cached) = self.cached_primary_read.clone() {
220            return Some(cached);
221        }
222
223        if offer.has_mime_type(&self.self_mime) {
224            return self.primary_contents.clone();
225        }
226
227        let item = offer
228            .read_text(&self.connection)
229            .or_else(|| offer.read_image(&self.connection))?;
230
231        self.cached_primary_read = Some(item.clone());
232        Some(item)
233    }
234
235    fn send_internal(&self, fd: OwnedFd, bytes: Vec<u8>) {
236        let mut written = 0;
237        self.loop_handle
238            .insert_source(
239                calloop::generic::Generic::new(
240                    File::from(fd),
241                    calloop::Interest::WRITE,
242                    calloop::Mode::Level,
243                ),
244                move |_, file, _| {
245                    let file = unsafe { file.get_mut() };
246                    loop {
247                        match file.write(&bytes[written..]) {
248                            Ok(n) if written + n == bytes.len() => {
249                                written += n;
250                                break Ok(PostAction::Remove);
251                            }
252                            Ok(n) => written += n,
253                            Err(err) if err.kind() == ErrorKind::WouldBlock => {
254                                break Ok(PostAction::Continue);
255                            }
256                            Err(_) => break Ok(PostAction::Remove),
257                        }
258                    }
259                },
260            )
261            .unwrap();
262    }
263}
264