22
33#![ allow( unsafe_code) ]
44
5- use std:: convert:: Infallible ;
6- use std:: env;
7- use std:: ffi:: { CStr , CString } ;
8- use std:: os:: fd:: { AsRawFd , RawFd } ;
9- use std:: pin:: Pin ;
10- use std:: task:: { Context , Poll } ;
11-
12- use anyhow:: Result ;
13- use close_fds:: CloseFdsBuilder ;
14- use nix:: errno:: Errno ;
15- use nix:: libc:: { login_tty, TIOCGWINSZ , TIOCSWINSZ } ;
16- use nix:: pty:: { self , Winsize } ;
17- use nix:: sys:: signal:: { kill, Signal :: SIGKILL } ;
18- use nix:: sys:: wait:: waitpid;
19- use nix:: unistd:: { execvp, fork, ForkResult , Pid } ;
20- use pin_project:: { pin_project, pinned_drop} ;
21- use tokio:: fs:: { self , File } ;
22- use tokio:: io:: { self , AsyncRead , AsyncWrite } ;
23- use tracing:: { instrument, trace} ;
24-
25- /// Returns the default shell on this system.
26- pub async fn get_default_shell ( ) -> String {
27- if let Ok ( shell) = env:: var ( "SHELL" ) {
28- if !shell. is_empty ( ) {
29- return shell;
30- }
31- }
32- for shell in [
33- "/bin/bash" ,
34- "/bin/sh" ,
35- "/usr/local/bin/bash" ,
36- "/usr/local/bin/sh" ,
37- ] {
38- if fs:: metadata ( shell) . await . is_ok ( ) {
39- return shell. to_string ( ) ;
40- }
41- }
42- String :: from ( "sh" )
43- }
44-
45- /// An object that stores the state for a terminal session.
46- #[ pin_project( PinnedDrop ) ]
47- pub struct Terminal {
48- child : Pid ,
49- #[ pin]
50- master_read : File ,
51- #[ pin]
52- master_write : File ,
53- }
54-
55- impl Terminal {
56- /// Create a new terminal, with attached PTY.
57- #[ instrument]
58- pub async fn new ( shell : & str ) -> Result < Terminal > {
59- let result = pty:: openpty ( None , None ) ?;
60-
61- // The slave file descriptor was created by openpty() and is forked here.
62- let child = Self :: fork_child ( shell, result. slave . as_raw_fd ( ) ) ?;
63-
64- // We need to clone the file object to prevent livelocks in Tokio, when multiple
65- // reads and writes happen concurrently on the same file descriptor. This is a
66- // current limitation of how the `tokio::fs::File` struct is implemented, due to
67- // its blocking I/O on a separate thread.
68- let master_read = File :: from ( std:: fs:: File :: from ( result. master ) ) ;
69- let master_write = master_read. try_clone ( ) . await ?;
70-
71- trace ! ( %child, "creating new terminal" ) ;
72-
73- Ok ( Self {
74- child,
75- master_read,
76- master_write,
77- } )
78- }
79-
80- /// Entry point for the child process, which spawns a shell.
81- fn fork_child ( shell : & str , slave_port : RawFd ) -> Result < Pid > {
82- let shell = CString :: new ( shell. to_owned ( ) ) ?;
83-
84- // Safety: This does not use any async-signal-unsafe operations in the child
85- // branch, such as memory allocation.
86- match unsafe { fork ( ) } ? {
87- ForkResult :: Parent { child } => Ok ( child) ,
88- ForkResult :: Child => match Self :: execv_child ( & shell, slave_port) {
89- Ok ( infallible) => match infallible { } ,
90- Err ( _) => std:: process:: exit ( 1 ) ,
91- } ,
92- }
93- }
94-
95- fn execv_child ( shell : & CStr , slave_port : RawFd ) -> Result < Infallible , Errno > {
96- // Safety: The slave file descriptor was created by openpty().
97- Errno :: result ( unsafe { login_tty ( slave_port) } ) ?;
98- // Safety: This is called immediately before an execv(), and there are no other
99- // threads in this process to interact with its file descriptor table.
100- unsafe { CloseFdsBuilder :: new ( ) . closefrom ( 3 ) } ;
101-
102- // Set terminal environment variables appropriately.
103- env:: set_var ( "TERM" , "xterm-256color" ) ;
104- env:: set_var ( "COLORTERM" , "truecolor" ) ;
105- env:: set_var ( "TERM_PROGRAM" , "sshx" ) ;
106- env:: remove_var ( "TERM_PROGRAM_VERSION" ) ;
107-
108- // Start the process.
109- execvp ( shell, & [ shell] )
110- }
111-
112- /// Get the window size of the TTY.
113- pub fn get_winsize ( & self ) -> Result < ( u16 , u16 ) > {
114- nix:: ioctl_read_bad!( ioctl_get_winsize, TIOCGWINSZ , Winsize ) ;
115- let mut winsize = make_winsize ( 0 , 0 ) ;
116- // Safety: The master file descriptor was created by openpty().
117- unsafe { ioctl_get_winsize ( self . master_read . as_raw_fd ( ) , & mut winsize) } ?;
118- Ok ( ( winsize. ws_row , winsize. ws_col ) )
119- }
120-
121- /// Set the window size of the TTY.
122- pub fn set_winsize ( & self , rows : u16 , cols : u16 ) -> Result < ( ) > {
123- nix:: ioctl_write_ptr_bad!( ioctl_set_winsize, TIOCSWINSZ , Winsize ) ;
124- let winsize = make_winsize ( rows, cols) ;
125- // Safety: The master file descriptor was created by openpty().
126- unsafe { ioctl_set_winsize ( self . master_read . as_raw_fd ( ) , & winsize) } ?;
127- Ok ( ( ) )
128- }
129- }
130-
131- // Redirect terminal reads to the read file object.
132- impl AsyncRead for Terminal {
133- fn poll_read (
134- self : Pin < & mut Self > ,
135- cx : & mut Context < ' _ > ,
136- buf : & mut io:: ReadBuf < ' _ > ,
137- ) -> Poll < io:: Result < ( ) > > {
138- self . project ( ) . master_read . poll_read ( cx, buf)
139- }
140- }
141-
142- // Redirect terminal writes to the write file object.
143- impl AsyncWrite for Terminal {
144- fn poll_write (
145- self : Pin < & mut Self > ,
146- cx : & mut Context < ' _ > ,
147- buf : & [ u8 ] ,
148- ) -> Poll < io:: Result < usize > > {
149- self . project ( ) . master_write . poll_write ( cx, buf)
150- }
151-
152- fn poll_flush ( self : Pin < & mut Self > , cx : & mut Context < ' _ > ) -> Poll < io:: Result < ( ) > > {
153- self . project ( ) . master_write . poll_flush ( cx)
154- }
155-
156- fn poll_shutdown ( self : Pin < & mut Self > , cx : & mut Context < ' _ > ) -> Poll < io:: Result < ( ) > > {
157- self . project ( ) . master_write . poll_shutdown ( cx)
158- }
159- }
160-
161- #[ pinned_drop]
162- impl PinnedDrop for Terminal {
163- fn drop ( self : Pin < & mut Self > ) {
164- let this = self . project ( ) ;
165- let child = * this. child ;
166- trace ! ( %child, "dropping terminal" ) ;
167-
168- // Kill the child process on closure so that it doesn't keep running.
169- kill ( child, SIGKILL ) . ok ( ) ;
170-
171- // Reap the zombie process in a background thread.
172- std:: thread:: spawn ( move || {
173- waitpid ( child, None ) . ok ( ) ;
174- } ) ;
175- }
176- }
177-
178- fn make_winsize ( rows : u16 , cols : u16 ) -> Winsize {
179- Winsize {
180- ws_row : rows,
181- ws_col : cols,
182- ws_xpixel : 0 , // ignored
183- ws_ypixel : 0 , // ignored
5+ cfg_if:: cfg_if! {
6+ if #[ cfg( unix) ] {
7+ mod unix;
8+ pub use unix:: { get_default_shell, Terminal } ;
9+ } else if #[ cfg( windows) ] {
10+ mod windows;
11+ pub use windows:: { get_default_shell, Terminal } ;
12+ } else {
13+ compile_error!( "unsupported platform for terminal driver" ) ;
18414 }
18515}
18616
@@ -192,7 +22,8 @@ mod tests {
19222
19323 #[ tokio:: test]
19424 async fn winsize ( ) -> Result < ( ) > {
195- let terminal = Terminal :: new ( "/bin/sh" ) . await ?;
25+ let shell = if cfg ! ( unix) { "/bin/sh" } else { "cmd.exe" } ;
26+ let mut terminal = Terminal :: new ( shell) . await ?;
19627 assert_eq ! ( terminal. get_winsize( ) ?, ( 0 , 0 ) ) ;
19728 terminal. set_winsize ( 120 , 72 ) ?;
19829 assert_eq ! ( terminal. get_winsize( ) ?, ( 120 , 72 ) ) ;
0 commit comments