1use std::sync::Arc;
2
3use bytes::Bytes;
4use tokio::io::{AsyncRead, AsyncWrite};
5use tokio::sync::mpsc::{Receiver, Sender};
6use tokio::sync::{Mutex, Notify};
7
8use crate::{ChannelId, ChannelOpenFailure, Error, Pty, Sig};
9
10pub mod io;
11
12mod channel_ref;
13pub use channel_ref::ChannelRef;
14
15mod channel_stream;
16pub use channel_stream::ChannelStream;
17
18#[derive(Debug)]
19#[non_exhaustive]
20pub enum ChannelMsg {
22 Open {
23 id: ChannelId,
24 max_packet_size: u32,
25 window_size: u32,
26 },
27 Data {
28 data: Bytes,
29 },
30 ExtendedData {
31 data: Bytes,
32 ext: u32,
33 },
34 Eof,
35 Close,
36 RequestPty {
38 want_reply: bool,
39 term: String,
40 col_width: u32,
41 row_height: u32,
42 pix_width: u32,
43 pix_height: u32,
44 terminal_modes: Vec<(Pty, u32)>,
45 },
46 RequestShell {
48 want_reply: bool,
49 },
50 Exec {
52 want_reply: bool,
53 command: Vec<u8>,
54 },
55 Signal {
57 signal: Sig,
58 },
59 RequestSubsystem {
61 want_reply: bool,
62 name: String,
63 },
64 RequestX11 {
66 want_reply: bool,
67 single_connection: bool,
68 x11_authentication_protocol: String,
69 x11_authentication_cookie: String,
70 x11_screen_number: u32,
71 },
72 SetEnv {
74 want_reply: bool,
75 variable_name: String,
76 variable_value: String,
77 },
78 WindowChange {
80 col_width: u32,
81 row_height: u32,
82 pix_width: u32,
83 pix_height: u32,
84 },
85 AgentForward {
87 want_reply: bool,
88 },
89
90 XonXoff {
92 client_can_do: bool,
93 },
94 ExitStatus {
96 exit_status: u32,
97 },
98 ExitSignal {
100 signal_name: Sig,
101 core_dumped: bool,
102 error_message: String,
103 lang_tag: String,
104 },
105 WindowAdjusted {
107 new_size: u32,
108 },
109 Success,
111 Failure,
113 OpenFailure(ChannelOpenFailure),
114}
115
116#[derive(Clone, Debug)]
117pub(crate) struct WindowSizeRef {
118 value: Arc<Mutex<u32>>,
119 notifier: Arc<Notify>,
120}
121
122impl WindowSizeRef {
123 pub(crate) fn new(initial: u32) -> Self {
124 let notifier = Arc::new(Notify::new());
125 Self {
126 value: Arc::new(Mutex::new(initial)),
127 notifier,
128 }
129 }
130
131 pub(crate) async fn update(&self, value: u32) {
132 *self.value.lock().await = value;
133 self.notifier.notify_one();
134 }
135
136 pub(crate) fn subscribe(&self) -> Arc<Notify> {
137 Arc::clone(&self.notifier)
138 }
139}
140
141pub struct ChannelReadHalf {
145 pub(crate) receiver: Receiver<ChannelMsg>,
146}
147
148impl std::fmt::Debug for ChannelReadHalf {
149 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
150 f.debug_struct("ChannelReadHalf").finish()
151 }
152}
153
154impl ChannelReadHalf {
155 pub async fn wait(&mut self) -> Option<ChannelMsg> {
157 self.receiver.recv().await
158 }
159
160 pub fn make_reader(&mut self) -> impl AsyncRead + '_ {
163 self.make_reader_ext(None)
164 }
165
166 pub fn make_reader_ext(&mut self, ext: Option<u32>) -> impl AsyncRead + '_ {
169 io::ChannelRx::new(self, ext)
170 }
171}
172
173pub struct ChannelWriteHalf<Send: From<(ChannelId, ChannelMsg)>> {
177 pub(crate) id: ChannelId,
178 pub(crate) sender: Sender<Send>,
179 pub(crate) max_packet_size: u32,
180 pub(crate) window_size: WindowSizeRef,
181}
182
183impl<S: From<(ChannelId, ChannelMsg)>> std::fmt::Debug for ChannelWriteHalf<S> {
184 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
185 f.debug_struct("ChannelWriteHalf")
186 .field("id", &self.id)
187 .finish()
188 }
189}
190
191impl<S: From<(ChannelId, ChannelMsg)> + Send + Sync + 'static> ChannelWriteHalf<S> {
192 pub async fn writable_packet_size(&self) -> usize {
195 self.max_packet_size
196 .min(*self.window_size.value.lock().await) as usize
197 }
198
199 pub fn id(&self) -> ChannelId {
200 self.id
201 }
202
203 #[allow(clippy::too_many_arguments)] pub async fn request_pty(
206 &self,
207 want_reply: bool,
208 term: &str,
209 col_width: u32,
210 row_height: u32,
211 pix_width: u32,
212 pix_height: u32,
213 terminal_modes: &[(Pty, u32)],
214 ) -> Result<(), Error> {
215 self.send_msg(ChannelMsg::RequestPty {
216 want_reply,
217 term: term.to_string(),
218 col_width,
219 row_height,
220 pix_width,
221 pix_height,
222 terminal_modes: terminal_modes.to_vec(),
223 })
224 .await
225 }
226
227 pub async fn request_shell(&self, want_reply: bool) -> Result<(), Error> {
229 self.send_msg(ChannelMsg::RequestShell { want_reply }).await
230 }
231
232 pub async fn exec<A: Into<Vec<u8>>>(&self, want_reply: bool, command: A) -> Result<(), Error> {
236 self.send_msg(ChannelMsg::Exec {
237 want_reply,
238 command: command.into(),
239 })
240 .await
241 }
242
243 pub async fn signal(&self, signal: Sig) -> Result<(), Error> {
245 self.send_msg(ChannelMsg::Signal { signal }).await
246 }
247
248 pub async fn request_subsystem<A: Into<String>>(
250 &self,
251 want_reply: bool,
252 name: A,
253 ) -> Result<(), Error> {
254 self.send_msg(ChannelMsg::RequestSubsystem {
255 want_reply,
256 name: name.into(),
257 })
258 .await
259 }
260
261 pub async fn request_x11<A: Into<String>, B: Into<String>>(
266 &self,
267 want_reply: bool,
268 single_connection: bool,
269 x11_authentication_protocol: A,
270 x11_authentication_cookie: B,
271 x11_screen_number: u32,
272 ) -> Result<(), Error> {
273 self.send_msg(ChannelMsg::RequestX11 {
274 want_reply,
275 single_connection,
276 x11_authentication_protocol: x11_authentication_protocol.into(),
277 x11_authentication_cookie: x11_authentication_cookie.into(),
278 x11_screen_number,
279 })
280 .await
281 }
282
283 pub async fn set_env<A: Into<String>, B: Into<String>>(
285 &self,
286 want_reply: bool,
287 variable_name: A,
288 variable_value: B,
289 ) -> Result<(), Error> {
290 self.send_msg(ChannelMsg::SetEnv {
291 want_reply,
292 variable_name: variable_name.into(),
293 variable_value: variable_value.into(),
294 })
295 .await
296 }
297
298 pub async fn window_change(
300 &self,
301 col_width: u32,
302 row_height: u32,
303 pix_width: u32,
304 pix_height: u32,
305 ) -> Result<(), Error> {
306 self.send_msg(ChannelMsg::WindowChange {
307 col_width,
308 row_height,
309 pix_width,
310 pix_height,
311 })
312 .await
313 }
314
315 pub async fn agent_forward(&self, want_reply: bool) -> Result<(), Error> {
317 self.send_msg(ChannelMsg::AgentForward { want_reply }).await
318 }
319
320 pub async fn data<R: tokio::io::AsyncRead + Unpin>(&self, data: R) -> Result<(), Error> {
322 self.send_data(None, data).await
323 }
324
325 pub async fn data_bytes(&self, data: impl Into<Bytes>) -> Result<(), Error> {
327 self.send_bytes(None, data.into()).await
328 }
329
330 pub async fn extended_data<R: tokio::io::AsyncRead + Unpin>(
334 &self,
335 ext: u32,
336 data: R,
337 ) -> Result<(), Error> {
338 self.send_data(Some(ext), data).await
339 }
340
341 pub async fn extended_data_bytes(
343 &self,
344 ext: u32,
345 data: impl Into<Bytes>,
346 ) -> Result<(), Error> {
347 self.send_bytes(Some(ext), data.into()).await
348 }
349
350 async fn send_data<R: tokio::io::AsyncRead + Unpin>(
351 &self,
352 ext: Option<u32>,
353 mut data: R,
354 ) -> Result<(), Error> {
355 let mut tx = self.make_writer_ext(ext);
356
357 tokio::io::copy(&mut data, &mut tx).await?;
358
359 Ok(())
360 }
361
362 async fn reserve_writable_chunk(&self, remaining: usize) -> Result<usize, Error> {
363 if self.max_packet_size == 0 {
364 return Err(Error::Inconsistent);
365 }
366 loop {
367 let mut window_size = self.window_size.value.lock().await;
368 let writable = (self.max_packet_size as usize)
369 .min(*window_size as usize)
370 .min(remaining);
371 if writable > 0 {
372 *window_size -= writable as u32;
373 if *window_size > 0 {
374 self.window_size.notifier.notify_one();
375 }
376 return Ok(writable);
377 }
378 let notified = self.window_size.notifier.notified();
379 drop(window_size);
380 notified.await;
381 }
382 }
383
384 async fn send_bytes(&self, ext: Option<u32>, data: Bytes) -> Result<(), Error> {
385 if data.is_empty() {
386 return Ok(());
387 }
388
389 let mut offset = 0;
390 while offset < data.len() {
391 let writable = self.reserve_writable_chunk(data.len() - offset).await?;
392 let end = offset + writable;
393 let chunk = data.slice(offset..end);
394 let msg = match ext {
395 None => ChannelMsg::Data { data: chunk },
396 Some(ext) => ChannelMsg::ExtendedData { data: chunk, ext },
397 };
398 self.send_msg(msg).await?;
399 offset = end;
400 }
401
402 Ok(())
403 }
404
405 pub async fn eof(&self) -> Result<(), Error> {
406 self.send_msg(ChannelMsg::Eof).await
407 }
408
409 pub async fn exit_status(&self, exit_status: u32) -> Result<(), Error> {
410 self.send_msg(ChannelMsg::ExitStatus { exit_status }).await
411 }
412
413 pub async fn close(&self) -> Result<(), Error> {
415 self.send_msg(ChannelMsg::Close).await
416 }
417
418 async fn send_msg(&self, msg: ChannelMsg) -> Result<(), Error> {
419 self.sender
420 .send((self.id, msg).into())
421 .await
422 .map_err(|_| Error::SendError)
423 }
424
425 pub fn make_writer(&self) -> impl AsyncWrite + 'static {
428 self.make_writer_ext(None)
429 }
430
431 pub fn make_writer_ext(&self, ext: Option<u32>) -> impl AsyncWrite + 'static {
434 io::ChannelTx::new(
435 self.sender.clone(),
436 self.id,
437 self.window_size.value.clone(),
438 self.window_size.subscribe(),
439 self.max_packet_size,
440 ext,
441 )
442 }
443}
444
445pub struct Channel<Send: From<(ChannelId, ChannelMsg)>> {
449 pub(crate) read_half: ChannelReadHalf,
450 pub(crate) write_half: ChannelWriteHalf<Send>,
451}
452
453impl<T: From<(ChannelId, ChannelMsg)>> std::fmt::Debug for Channel<T> {
454 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
455 f.debug_struct("Channel")
456 .field("id", &self.write_half.id)
457 .finish()
458 }
459}
460
461impl<S: From<(ChannelId, ChannelMsg)> + Send + Sync + 'static> Channel<S> {
462 pub(crate) fn new(
463 id: ChannelId,
464 sender: Sender<S>,
465 max_packet_size: u32,
466 window_size: u32,
467 channel_buffer_size: usize,
468 ) -> (Self, ChannelRef) {
469 let (tx, rx) = tokio::sync::mpsc::channel(channel_buffer_size);
470 let window_size = WindowSizeRef::new(window_size);
471 let read_half = ChannelReadHalf { receiver: rx };
472 let write_half = ChannelWriteHalf {
473 id,
474 sender,
475 max_packet_size,
476 window_size: window_size.clone(),
477 };
478
479 (
480 Self {
481 write_half,
482 read_half,
483 },
484 ChannelRef {
485 sender: tx,
486 window_size,
487 },
488 )
489 }
490
491 pub async fn writable_packet_size(&self) -> usize {
494 self.write_half.writable_packet_size().await
495 }
496
497 pub fn id(&self) -> ChannelId {
498 self.write_half.id()
499 }
500
501 pub fn split(self) -> (ChannelReadHalf, ChannelWriteHalf<S>) {
504 (self.read_half, self.write_half)
505 }
506
507 #[allow(clippy::too_many_arguments)] pub async fn request_pty(
510 &self,
511 want_reply: bool,
512 term: &str,
513 col_width: u32,
514 row_height: u32,
515 pix_width: u32,
516 pix_height: u32,
517 terminal_modes: &[(Pty, u32)],
518 ) -> Result<(), Error> {
519 self.write_half
520 .request_pty(
521 want_reply,
522 term,
523 col_width,
524 row_height,
525 pix_width,
526 pix_height,
527 terminal_modes,
528 )
529 .await
530 }
531
532 pub async fn request_shell(&self, want_reply: bool) -> Result<(), Error> {
534 self.write_half.request_shell(want_reply).await
535 }
536
537 pub async fn exec<A: Into<Vec<u8>>>(&self, want_reply: bool, command: A) -> Result<(), Error> {
541 self.write_half.exec(want_reply, command).await
542 }
543
544 pub async fn signal(&self, signal: Sig) -> Result<(), Error> {
546 self.write_half.signal(signal).await
547 }
548
549 pub async fn request_subsystem<A: Into<String>>(
551 &self,
552 want_reply: bool,
553 name: A,
554 ) -> Result<(), Error> {
555 self.write_half.request_subsystem(want_reply, name).await
556 }
557
558 pub async fn request_x11<A: Into<String>, B: Into<String>>(
563 &self,
564 want_reply: bool,
565 single_connection: bool,
566 x11_authentication_protocol: A,
567 x11_authentication_cookie: B,
568 x11_screen_number: u32,
569 ) -> Result<(), Error> {
570 self.write_half
571 .request_x11(
572 want_reply,
573 single_connection,
574 x11_authentication_protocol,
575 x11_authentication_cookie,
576 x11_screen_number,
577 )
578 .await
579 }
580
581 pub async fn set_env<A: Into<String>, B: Into<String>>(
583 &self,
584 want_reply: bool,
585 variable_name: A,
586 variable_value: B,
587 ) -> Result<(), Error> {
588 self.write_half
589 .set_env(want_reply, variable_name, variable_value)
590 .await
591 }
592
593 pub async fn window_change(
595 &self,
596 col_width: u32,
597 row_height: u32,
598 pix_width: u32,
599 pix_height: u32,
600 ) -> Result<(), Error> {
601 self.write_half
602 .window_change(col_width, row_height, pix_width, pix_height)
603 .await
604 }
605
606 pub async fn agent_forward(&self, want_reply: bool) -> Result<(), Error> {
608 self.write_half.agent_forward(want_reply).await
609 }
610
611 pub async fn data<R: tokio::io::AsyncRead + Unpin>(&self, data: R) -> Result<(), Error> {
613 self.write_half.data(data).await
614 }
615
616 pub async fn data_bytes(&self, data: impl Into<Bytes>) -> Result<(), Error> {
618 self.write_half.data_bytes(data).await
619 }
620
621 pub async fn extended_data<R: tokio::io::AsyncRead + Unpin>(
625 &self,
626 ext: u32,
627 data: R,
628 ) -> Result<(), Error> {
629 self.write_half.extended_data(ext, data).await
630 }
631
632 pub async fn extended_data_bytes(
634 &self,
635 ext: u32,
636 data: impl Into<Bytes>,
637 ) -> Result<(), Error> {
638 self.write_half.extended_data_bytes(ext, data).await
639 }
640
641 pub async fn eof(&self) -> Result<(), Error> {
642 self.write_half.eof().await
643 }
644
645 pub async fn exit_status(&self, exit_status: u32) -> Result<(), Error> {
646 self.write_half.exit_status(exit_status).await
647 }
648
649 pub async fn close(&self) -> Result<(), Error> {
651 self.write_half.close().await
652 }
653
654 pub async fn wait(&mut self) -> Option<ChannelMsg> {
656 self.read_half.wait().await
657 }
658
659 pub fn into_stream(self) -> ChannelStream<S> {
662 ChannelStream::new(
663 io::ChannelTx::new(
664 self.write_half.sender.clone(),
665 self.write_half.id,
666 self.write_half.window_size.value.clone(),
667 self.write_half.window_size.subscribe(),
668 self.write_half.max_packet_size,
669 None,
670 ),
671 io::ChannelRx::new(io::ChannelCloseOnDrop(self), None),
672 )
673 }
674
675 pub fn make_reader(&mut self) -> impl AsyncRead + '_ {
678 self.read_half.make_reader()
679 }
680
681 pub fn make_reader_ext(&mut self, ext: Option<u32>) -> impl AsyncRead + '_ {
684 self.read_half.make_reader_ext(ext)
685 }
686
687 pub fn make_writer(&self) -> impl AsyncWrite + 'static {
690 self.write_half.make_writer()
691 }
692
693 pub fn make_writer_ext(&self, ext: Option<u32>) -> impl AsyncWrite + 'static {
696 self.write_half.make_writer_ext(ext)
697 }
698}
699
700#[cfg(test)]
701mod tests {
702 use tokio::sync::mpsc;
703
704 use super::*;
705
706 fn test_write_half(
707 window_size: WindowSizeRef,
708 max_packet_size: u32,
709 ) -> (
710 ChannelWriteHalf<(ChannelId, ChannelMsg)>,
711 mpsc::Receiver<(ChannelId, ChannelMsg)>,
712 ) {
713 let (sender, receiver) = mpsc::channel(8);
714 (
715 ChannelWriteHalf {
716 id: ChannelId(7),
717 sender,
718 max_packet_size,
719 window_size,
720 },
721 receiver,
722 )
723 }
724
725 #[tokio::test]
726 async fn data_bytes_sends_one_owned_message_when_window_permits() {
727 let payload = Bytes::from_static(b"owned data");
728 let (write_half, mut receiver) = test_write_half(WindowSizeRef::new(1024), 1024);
729
730 write_half.data_bytes(payload.clone()).await.unwrap();
731
732 match receiver.recv().await.unwrap() {
733 (ChannelId(7), ChannelMsg::Data { data }) => {
734 assert_eq!(data, payload);
735 assert_eq!(data.as_ptr(), payload.as_ptr());
736 }
737 msg => panic!("unexpected message: {msg:?}"),
738 }
739 }
740
741 #[tokio::test]
742 async fn data_bytes_splits_by_max_packet_size_without_copying() {
743 let payload = Bytes::from_static(b"abcdefghij");
744 let (write_half, mut receiver) = test_write_half(WindowSizeRef::new(1024), 4);
745
746 write_half.data_bytes(payload.clone()).await.unwrap();
747
748 for (range, expected) in [
749 (0..4, &b"abcd"[..]),
750 (4..8, &b"efgh"[..]),
751 (8..10, &b"ij"[..]),
752 ] {
753 match receiver.recv().await.unwrap() {
754 (ChannelId(7), ChannelMsg::Data { data }) => {
755 assert_eq!(data.as_ref(), expected);
756 assert_eq!(data.as_ptr(), payload.slice(range).as_ptr());
757 }
758 msg => panic!("unexpected message: {msg:?}"),
759 }
760 }
761 assert!(receiver.try_recv().is_err());
762 }
763
764 #[tokio::test]
765 async fn extended_data_bytes_preserves_extension_code() {
766 let payload = Bytes::from_static(b"stderr");
767 let (write_half, mut receiver) = test_write_half(WindowSizeRef::new(1024), 1024);
768
769 write_half
770 .extended_data_bytes(1, payload.clone())
771 .await
772 .unwrap();
773
774 match receiver.recv().await.unwrap() {
775 (ChannelId(7), ChannelMsg::ExtendedData { data, ext }) => {
776 assert_eq!(ext, 1);
777 assert_eq!(data, payload);
778 assert_eq!(data.as_ptr(), payload.as_ptr());
779 }
780 msg => panic!("unexpected message: {msg:?}"),
781 }
782 }
783
784 #[tokio::test]
785 async fn data_bytes_empty_payload_sends_nothing() {
786 let (write_half, mut receiver) = test_write_half(WindowSizeRef::new(1024), 1024);
787
788 write_half.data_bytes(Bytes::new()).await.unwrap();
789
790 assert!(receiver.try_recv().is_err());
791 }
792
793 #[tokio::test]
794 async fn data_bytes_waits_for_window_update() {
795 let window_size = WindowSizeRef::new(0);
796 let (write_half, mut receiver) = test_write_half(window_size.clone(), 1024);
797 let send = tokio::spawn(async move {
798 write_half
799 .data_bytes(Bytes::from_static(b"after-window"))
800 .await
801 .unwrap();
802 });
803
804 tokio::task::yield_now().await;
805 assert!(!send.is_finished());
806
807 window_size.update(1024).await;
808 send.await.unwrap();
809
810 match receiver.recv().await.unwrap() {
811 (ChannelId(7), ChannelMsg::Data { data }) => {
812 assert_eq!(data.as_ref(), b"after-window");
813 }
814 msg => panic!("unexpected message: {msg:?}"),
815 }
816 }
817
818 #[tokio::test]
819 async fn data_bytes_rejects_zero_max_packet_size() {
820 let (write_half, mut receiver) = test_write_half(WindowSizeRef::new(1024), 0);
821
822 let result = write_half.data_bytes(Bytes::from_static(b"owned")).await;
823
824 assert!(matches!(result, Err(Error::Inconsistent)));
825 assert!(receiver.try_recv().is_err());
826 }
827
828 #[tokio::test]
829 async fn channel_data_bytes_forwards_to_write_half() {
830 let (sender, mut receiver) = mpsc::channel(8);
831 let (channel, _reference) =
832 Channel::<(ChannelId, ChannelMsg)>::new(ChannelId(9), sender, 1024, 1024, 8);
833
834 channel.data_bytes(Bytes::from_static(b"channel")).await.unwrap();
835
836 match receiver.recv().await.unwrap() {
837 (ChannelId(9), ChannelMsg::Data { data }) => {
838 assert_eq!(data.as_ref(), b"channel");
839 }
840 msg => panic!("unexpected message: {msg:?}"),
841 }
842 }
843}