pyc_shell/shell/proc/
process.rs

1//! ## Process
2//!
3//! `Process` contains the implementation for ShellProc
4
5/*
6*
7*   Copyright (C) 2020 Christian Visintin - christian.visintin1997@gmail.com
8*
9* 	This file is part of "Pyc"
10*
11*   Pyc is free software: you can redistribute it and/or modify
12*   it under the terms of the GNU General Public License as published by
13*   the Free Software Foundation, either version 3 of the License, or
14*   (at your option) any later version.
15*
16*   Pyc is distributed in the hope that it will be useful,
17*   but WITHOUT ANY WARRANTY; without even the implied warranty of
18*   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
19*   GNU General Public License for more details.
20*
21*   You should have received a copy of the GNU General Public License
22*   along with Pyc.  If not, see <http://www.gnu.org/licenses/>.
23*
24*/
25
26extern crate nix;
27extern crate tempfile;
28extern crate uuid;
29
30use super::{ShellError, ShellProc, ShellProcState};
31use super::pipe::Pipe;
32
33use std::ffi::{CStr, CString};
34use std::os::unix::io::RawFd;
35use std::path::PathBuf;
36use std::time::{Duration, Instant};
37use uuid::Uuid;
38
39impl ShellProc {
40
41    /// ### start
42    /// 
43    /// Start a process
44    pub fn start(argv: Vec<String>) -> Result<ShellProc, ShellError> {
45        if argv.len() == 0 {
46            return Err(ShellError::CouldNotStartProcess)
47        }
48        //Generate UUID - NOTE: UUID is used to notice process that shell subprocess has terminated
49        let uuid: String = Uuid::new_v4().to_hyphenated().to_string();
50        //Create pipes
51        let tmpdir: tempfile::TempDir = tempfile::TempDir::new().unwrap();
52        let stdin_pipe: Pipe = match Pipe::open(&tmpdir.path().join("stdin.fifo")) {
53            Ok(p) => p,
54            Err(err) => return Err(err)
55        };
56        let stderr_pipe: Pipe = match Pipe::open(&tmpdir.path().join("stderr.fifo")) {
57            Ok(p) => p,
58            Err(err) => return Err(err)
59        };
60        let stdout_pipe: Pipe = match Pipe::open(&tmpdir.path().join("stdout.fifo")) {
61            Ok(p) => p,
62            Err(err) => return Err(err)
63        };
64        //Fork process
65        match unsafe {nix::unistd::fork()} {
66            Ok(nix::unistd::ForkResult::Parent { child, .. }) => {
67                //Prepare echo command
68                //FIXME: handle fish $status
69                let echo_command: String = format!("echo \"\x02$?;`pwd`;{}\x03\"\n", uuid);
70                let wrkdir: PathBuf = match std::env::current_dir() {
71                    Err(_) => PathBuf::from("/"),
72                    Ok(path) => PathBuf::from(path.as_path())
73                };
74                //Return Shell Proc
75                Ok(ShellProc {
76                    state: ShellProcState::Idle,
77                    uuid: uuid,
78                    exit_status: 0,
79                    exec_time: Duration::from_millis(0),
80                    wrkdir: wrkdir,
81                    pid: child.as_raw(),
82                    rc: 255,
83                    stdout_cache: None,
84                    start_time: Instant::now(),
85                    echo_command: echo_command,
86                    stdin_pipe: stdin_pipe,
87                    stderr_pipe: stderr_pipe,
88                    stdout_pipe: stdout_pipe
89                })
90            },
91            Ok(nix::unistd::ForkResult::Child) => {
92                std::process::exit(ShellProc::run(argv, stdin_pipe.fd, stderr_pipe.fd, stdout_pipe.fd));
93            },
94            Err(_) => {
95                return Err(ShellError::CouldNotStartProcess)
96            }
97        }
98    }
99
100    /// ### cleanup
101    /// 
102    /// cleanup shell once exited. Returns the shell exit code
103    pub fn cleanup(&mut self) -> Result<u8, ShellError> {
104        if self.update_state() != ShellProcState::Terminated {
105            return Err(ShellError::ShellRunning)
106        }
107        //Close pipes
108        let _ = self.stdin_pipe.close();
109        let _ = self.stdout_pipe.close();
110        let _ = self.stderr_pipe.close();
111        Ok(self.rc)
112    }
113
114    /// ### raise
115    /// 
116    /// Send signal to shell
117    pub fn raise(&self, signal: nix::sys::signal::Signal) -> Result<(), ShellError> {
118        match nix::sys::signal::kill(nix::unistd::Pid::from_raw(self.pid), signal) {
119            Ok(_) => Ok(()),
120            Err(_) => Err(ShellError::CouldNotKill)
121        }
122    }
123
124    /// ### kill
125    /// 
126    /// Kill shell sending SIGKILL
127    pub fn kill(&self) -> Result<(), ShellError> {
128        self.raise(nix::sys::signal::Signal::SIGKILL)
129    }
130    
131    /// ### read
132    /// 
133    /// Read from child pipes
134    pub fn read(&mut self) -> Result<(Option<String>, Option<String>), ShellError> {
135        /* NOTE: doesn't make sense; read must be possible even if shell has terminated
136        if self.update_state() == ShellProcState::Terminated {
137            return Err(ShellError::ShellTerminated)
138        }*/
139        let stdout: Option<String> = match self.stdout_pipe.read(50, false) {
140            Ok(stdout) => self.parse_stdout(stdout),
141            Err(err) => return Err(err)
142        };
143        let stderr: Option<String> = match self.stderr_pipe.read(50, false) {
144            Ok(stderr) => match stderr {
145                None => None,
146                Some(stderr) => Some(stderr)
147            },
148            Err(err) => return Err(err)
149        };
150        Ok((stdout, stderr))
151    }
152
153    /// ### write
154    /// 
155    /// Write to child process stdin
156    pub fn write(&mut self, mut data: String) -> Result<(), ShellError> {
157        if self.update_state() == ShellProcState::Terminated {
158            return Err(ShellError::ShellTerminated)
159        }
160        //Add echo command to data if shell state is Idle
161        if self.state == ShellProcState::Idle {
162            //Replace data newline with ';'
163            while data.ends_with('\n') {
164                data.pop();
165            }
166            //Append semicolon to data
167            if ! data.ends_with(';') {
168                data.push(';');
169            }
170            //Append echo command to data
171            data.push_str(self.echo_command.as_str());
172            //Set state to running
173            self.set_state_running();
174        }
175        self.stdin_pipe.write(data, 5000)
176    }
177
178    /// ### run
179    /// 
180    /// Run method for thread
181    fn run(argv: Vec<String>, stdin: RawFd, stderr: RawFd, stdout: RawFd) -> i32 {
182        //Set child process stdout/stdin/stderr
183        if let Err(_) = nix::unistd::dup2(stdin, 0) {
184            return 255
185        }
186        if let Err(_) = nix::unistd::dup2(stdout, 1) {
187            return 255
188        }
189        if let Err(_) = nix::unistd::dup2(stderr, 2) {
190            return 255
191        }
192        //Prepare arguments
193        let mut c_argv: Vec<CString> = Vec::with_capacity(argv.len());
194        for arg in argv.iter() {
195            c_argv.push(CString::new(arg.as_str()).unwrap());
196        }
197        let mut c_argv_refs: Vec<&CStr> = Vec::with_capacity(c_argv.len());
198        for arg in c_argv.iter() {
199            c_argv_refs.push(arg);
200        }
201        //Exec process
202        if let Err(_) = nix::unistd::execvp(c_argv_refs.get(0).unwrap(), c_argv_refs.as_slice()) {
203            return 255
204        }
205        return 0
206    }
207
208    /// ### update_state
209    /// 
210    /// Update shell running state checking if the other thread has terminated
211    pub fn update_state(&mut self) -> ShellProcState {
212        //Wait pid (NO HANG)
213        match nix::sys::wait::waitpid(nix::unistd::Pid::from_raw(self.pid), Some(nix::sys::wait::WaitPidFlag::WNOHANG)) {
214            Err(_) => {}, //Could not get information
215            Ok(status) => match status {
216                nix::sys::wait::WaitStatus::Exited(_, rc) => {
217                    self.state = ShellProcState::Terminated;
218                    self.rc = rc as u8;
219                },
220                nix::sys::wait::WaitStatus::Signaled(_, signal, _) => {
221                    self.state = ShellProcState::Terminated;
222                    self.rc = signal as u8;
223                },
224                _ => {}, //Still running
225            }
226        };
227        self.state
228    }
229
230    /// ### parse_stdout
231    /// 
232    /// Parse stdout received from shell process
233    fn parse_stdout(&mut self, stdout: Option<String>) -> Option<String> {
234        match stdout {
235            None => None,
236            Some(stdout) => {
237                //Treat stdout
238                let termination_string: String = format!("{}\x03\n", self.uuid);
239                //Check if ends with this (\x02${?};${PWD};${UUID}\x03\n)
240                //Create check string (cache + stdout)
241                let check_string: String = match &self.stdout_cache {
242                    None => stdout.clone(),
243                    Some(cache) => {
244                        let mut s: String = String::with_capacity(stdout.len() + cache.len());
245                        s.push_str(cache.as_str());
246                        s.push_str(stdout.as_str());
247                        s
248                    }
249                };
250                //Check if string ends with termination string
251                if check_string.ends_with(termination_string.as_str()) {
252                    //It's the end of shell execution, split string in output and METADATA
253                    //Let's find the index of \x02
254                    let mut stx_index: usize = check_string.len();
255                    for c in check_string.chars().rev() {
256                        if c == '\x02' {
257                            break;
258                        }
259                        stx_index -= 1; //Decrement STX index
260                    }
261                    let metadata: String = String::from(&check_string[stx_index..check_string.len() - 2]);
262                    //Get stdout
263                    let stx_index_stdout: usize = stx_index - match &self.stdout_cache {
264                        Some(s) => s.len(),
265                        None => 0
266                    };
267                    let stdout: String = String::from(&stdout[..stx_index_stdout - 1]);
268                    //get metadata
269                    self.set_state_idle(metadata);
270                    //Clear cache
271                    self.stdout_cache = None;
272                    match stdout.len() {
273                        0 => None,
274                        _ => Some(stdout)
275                    }
276                } else {
277                    //Not a termination
278                    //Push stdout to cache
279                    self.stdout_cache = Some(stdout.clone());
280                    //Return stdout
281                    Some(stdout)
282                }
283            }
284        }
285    }
286
287    /// ### set_state_idle
288    /// 
289    /// Parse metadata string and set state back to idle
290    fn set_state_idle(&mut self, metadata: String) {
291        for (index, token) in metadata.split(";").enumerate() {
292            match index {
293                0 => self.exit_status = token.parse::<u8>().unwrap_or(255),
294                1 => self.wrkdir = PathBuf::from(token),
295                _ => continue
296            }
297        }
298        self.exec_time = self.start_time.elapsed();
299        self.state = ShellProcState::Idle;
300    }
301
302    /// ### set_state_running
303    /// 
304    /// Set state to running
305    fn set_state_running(&mut self) {
306        self.start_time = Instant::now();
307        self.state = ShellProcState::SubprocessRunning;
308    }
309}
310
311impl Drop for ShellProc {
312    fn drop(&mut self) {
313        if let Err(_) = self.cleanup() {
314            let _ = self.kill(); //Force to terminate
315            let _ = self.cleanup(); //Then finally clean up
316        }
317    }
318}
319
320//@! Test module
321
322#[cfg(test)]
323mod tests {
324
325    use super::*;
326
327    use nix::NixPath;
328    use std::time::Duration;
329    use std::thread::sleep;
330
331    #[test]
332    fn test_process_start_stop() {
333        let mut shell_proc: ShellProc = ShellProc::start(vec![String::from("sh")]).unwrap();
334        println!("A new shell started with PID {}", shell_proc.pid);
335        //Check shell parameters
336        assert_eq!(shell_proc.state, ShellProcState::Idle);
337        assert_eq!(shell_proc.exit_status, 0);
338        assert_ne!(shell_proc.pid, 0);
339        assert_ne!(shell_proc.wrkdir.len(), 0);
340        assert_eq!(shell_proc.exec_time, Duration::from_millis(0));
341        assert_eq!(shell_proc.rc, 255);
342        assert_ne!(shell_proc.uuid.len(), 0);
343        assert!(shell_proc.stdout_cache.is_none());
344        assert_eq!(shell_proc.echo_command, format!("echo \"\x02$?;`pwd`;{}\x03\"\n", shell_proc.uuid));
345        //Verify shell is still running
346        sleep(Duration::from_millis(500));
347        assert_eq!(shell_proc.update_state(), ShellProcState::Idle);
348        //Stop process
349        assert!(shell_proc.kill().is_ok());
350        sleep(Duration::from_millis(500));
351        assert_eq!(shell_proc.update_state(), ShellProcState::Terminated);
352        //Rc should be set to 9
353        assert_eq!(shell_proc.state, ShellProcState::Terminated);
354        assert_eq!(shell_proc.rc, 9);
355        //Cleanup
356        assert!(shell_proc.cleanup().is_ok());
357    }
358
359    #[test]
360    fn test_process_start_error() {
361        let mut shell_proc: ShellProc = ShellProc::start(vec![String::from("piroporopero")]).unwrap();
362        println!("A new shell started with PID {}", shell_proc.pid);
363        //Shell should have died
364        sleep(Duration::from_millis(1000));
365        assert_eq!(shell_proc.update_state(), ShellProcState::Terminated);
366        assert_eq!(shell_proc.rc, 255);
367    }
368
369    #[test]
370    fn test_process_raise() {
371        let mut shell_proc: ShellProc = ShellProc::start(vec![String::from("sh")]).unwrap();
372        println!("A new shell started with PID {}", shell_proc.pid);
373        //Verify shell is still running
374        sleep(Duration::from_millis(500));
375        assert_eq!(shell_proc.update_state(), ShellProcState::Idle);
376        //Send SIGINT
377        assert!(shell_proc.raise(nix::sys::signal::Signal::SIGINT).is_ok());
378        sleep(Duration::from_millis(500));
379        assert_eq!(shell_proc.update_state(), ShellProcState::Terminated);
380        assert_eq!(shell_proc.rc, 2);
381    }
382
383    #[test]
384    fn test_process_parse_metadata() {
385        let mut shell_proc: ShellProc = ShellProc::start(vec![String::from("sh")]).unwrap();
386        println!("A new shell started with PID {}", shell_proc.pid);
387        sleep(Duration::from_millis(500)); //DON'T REMOVE THIS SLEEP
388        //Parse metadata
389        let metadata: String = String::from("128;/home;ee9ec814-a751-4329-850f-6d54d12c8a5c");
390        shell_proc.state = ShellProcState::SubprocessRunning;
391        shell_proc.set_state_idle(metadata);
392        //Verify metadata have been parsed successfully
393        assert_eq!(shell_proc.exit_status, 128);
394        assert_eq!(shell_proc.wrkdir, PathBuf::from("/home"));
395        assert_eq!(shell_proc.state, ShellProcState::Idle);
396        //Kill
397        assert!(shell_proc.kill().is_ok());
398    }
399
400    #[test]
401    fn test_process_parse_stdout() {
402        let mut shell_proc: ShellProc = ShellProc::start(vec![String::from("sh")]).unwrap();
403        println!("A new shell started with PID {}", shell_proc.pid);
404        sleep(Duration::from_millis(500)); //DON'T REMOVE THIS SLEEP
405        //Parse stdout when empty
406        assert!(shell_proc.parse_stdout(None).is_none());
407        //Parse stdout with metadata only (and parse theme)
408        shell_proc.state = ShellProcState::SubprocessRunning;
409        assert!(shell_proc.parse_stdout(Some(format!("\x02128;/home;{}\x03\n", shell_proc.uuid))).is_none());
410        assert_eq!(shell_proc.exit_status, 128);
411        assert_eq!(shell_proc.wrkdir, PathBuf::from("/home"));
412        assert_eq!(shell_proc.state, ShellProcState::Idle);
413        //Parse stdout with output only
414        shell_proc.state = ShellProcState::SubprocessRunning;
415        assert_eq!(shell_proc.parse_stdout(Some(String::from("HELLO\n"))).unwrap(), String::from("HELLO\n"));
416        assert_eq!(shell_proc.state, ShellProcState::SubprocessRunning); //State unchanged
417        assert_eq!(*shell_proc.stdout_cache.as_ref().unwrap(), String::from("HELLO\n"));
418        //Parse stdout with everything
419        shell_proc.state = ShellProcState::SubprocessRunning;
420        assert_eq!(shell_proc.parse_stdout(Some(format!("HELLO\n\x022;/tmp;{}\x03\n", shell_proc.uuid))).unwrap(), String::from("HELLO\n"));
421        assert_eq!(shell_proc.exit_status, 2);
422        assert_eq!(shell_proc.wrkdir, PathBuf::from("/tmp"));
423        assert_eq!(shell_proc.state, ShellProcState::Idle);
424        assert!(shell_proc.stdout_cache.is_none());
425        //Kill
426        assert!(shell_proc.kill().is_ok());
427    }
428
429    #[test]
430    fn test_process_command() {
431        let mut shell_proc: ShellProc = ShellProc::start(vec![String::from("sh")]).unwrap();
432        println!("A new shell started with PID {}", shell_proc.pid);
433        //Send a cd command
434        assert!(shell_proc.write(String::from("cd /tmp\n")).is_ok());
435        //State should have changed to subprocess
436        assert_eq!(shell_proc.state, ShellProcState::SubprocessRunning);
437        //Then read response
438        sleep(Duration::from_millis(50));
439        let (stdout, stderr) = shell_proc.read().unwrap();
440        //Output should be empty
441        assert!(stdout.is_none());
442        assert!(stderr.is_none());
443        //Verify shell is still running
444        sleep(Duration::from_millis(100));
445        assert_eq!(shell_proc.update_state(), ShellProcState::Idle);
446        //Verify wrkdir is now /tmp/
447        assert_eq!(shell_proc.wrkdir, PathBuf::from("/tmp"));
448        //Verify exit status
449        assert_eq!(shell_proc.exit_status, 0);
450        //Verify execution time
451        assert_ne!(shell_proc.exec_time.as_nanos(), 0);
452        //Stop process
453        assert!(shell_proc.kill().is_ok());
454        sleep(Duration::from_millis(500));
455        assert_eq!(shell_proc.update_state(), ShellProcState::Terminated);
456        //Rc should be set to 9
457        assert_eq!(shell_proc.state, ShellProcState::Terminated);
458        assert_eq!(shell_proc.rc, 9);
459        //Cleanup
460        assert!(shell_proc.cleanup().is_ok());
461    }
462
463}
pyc_shell/shell/proc/process.rs

pyc_shell/shell/proc/
process.rs
