/rain_server/src/start/ssh.rs

https://github.com/substantic/rain · Rust · 173 lines · 151 code · 17 blank · 5 comment · 8 complexity · 25a88cbc28ad77e14bf9bfc9c36eb6dd MD5 · raw file

  1. use std::error::Error;
  2. use std::io::BufRead;
  3. use std::io::Write;
  4. use std::path::Path;
  5. use std::process::{Command, Stdio};
  6. use error_chain::bail;
  7. use rain_core::errors::Result;
  8. use start::common::Readiness;
  9. use std::io::BufReader;
  10. pub struct RemoteProcess {
  11. name: String,
  12. host: String,
  13. pid: i32,
  14. readiness: Readiness,
  15. }
  16. impl RemoteProcess {
  17. pub fn new(name: String, host: &str, readiness: Readiness) -> Self {
  18. RemoteProcess {
  19. name,
  20. host: host.to_string(),
  21. pid: 0,
  22. readiness,
  23. }
  24. }
  25. fn create_ssh_command(&self) -> Command {
  26. let mut command = Command::new("ssh");
  27. command
  28. .arg("-o StrictHostKeyChecking=no")
  29. .arg(&self.host)
  30. .arg("/bin/sh")
  31. .stdin(Stdio::piped())
  32. .stdout(Stdio::piped())
  33. .stderr(Stdio::piped());
  34. command
  35. }
  36. pub fn run_ssh_first_line(&self, command: &str) -> Result<String> {
  37. let mut child = self.create_ssh_command()
  38. .spawn()
  39. .map_err(|e| format!("Start of 'ssh' failed: {}", e.description()))?;
  40. {
  41. let stdin = child.stdin.as_mut().unwrap();
  42. stdin.write_all(command.as_bytes())?;
  43. stdin.flush()?;
  44. }
  45. let mut output = String::new();
  46. {
  47. let mut reader = BufReader::new(child.stdout.as_mut().unwrap());
  48. reader.read_line(&mut output)?;
  49. }
  50. if output.is_empty() {
  51. let out = child.wait_with_output()?;
  52. bail!("ssh failed with: {}", ::std::str::from_utf8(&out.stderr)?);
  53. }
  54. Ok(output)
  55. }
  56. pub fn run_ssh(&self, command: &str) -> Result<(String, String)> {
  57. let mut child = self.create_ssh_command()
  58. .spawn()
  59. .map_err(|e| format!("Start of 'ssh' failed: {}", e.description()))?;
  60. {
  61. let stdin = child.stdin.as_mut().unwrap();
  62. stdin.write_all(command.as_bytes())?;
  63. stdin.flush()?;
  64. }
  65. let output = child.wait_with_output()?;
  66. let stderr = ::std::str::from_utf8(&output.stderr)?;
  67. let stdout = ::std::str::from_utf8(&output.stdout)?;
  68. if !output.status.success() {
  69. bail!("Connection to remote site failed: {}", stderr);
  70. }
  71. Ok((stdout.to_string(), stderr.to_string()))
  72. }
  73. pub fn start(&mut self, command: &str, cwd: &Path, log_dir: &Path) -> Result<()> {
  74. /* Shell script that has following goals:
  75. - check that log files are accessible
  76. - start a new sub-shell with desired command
  77. - return PID of the new shell and then terminate
  78. */
  79. let log_out = log_dir.join(&format!("{}.out", self.name));
  80. let log_err = log_dir.join(&format!("{}.err", self.name));
  81. let shell_cmd = format!(
  82. "\n
  83. mkdir -p {log_dir:?} || (echo \"Error: Cannot create log directory\"; exit 1)\n
  84. touch {log_out:?} || (echo \"Error: Cannot create log file\"; exit 1)\n
  85. touch {log_err:?} || (echo \"Error: Cannot create log file\"; exit 1)\n
  86. ({{\n
  87. cd {cwd:?} || exit 1;\n
  88. {command}\n
  89. }} > {log_out:?} 2> {log_err:?})&\n
  90. echo \"Ok: $!\"\n
  91. ",
  92. cwd = cwd,
  93. command = command,
  94. log_dir = log_dir,
  95. log_out = log_out,
  96. log_err = log_err
  97. );
  98. let stdout = self.run_ssh_first_line(&shell_cmd)?;
  99. if stdout.starts_with("Ok: ") {
  100. self.pid = stdout[4..]
  101. .trim()
  102. .parse()
  103. .map_err(|_| format!("Internal error, value is not integer: {}", stdout))?;
  104. } else if stdout.starts_with("Error: ") {
  105. bail!(
  106. "Remote process at {}, the following error occurs: {}",
  107. self.host,
  108. &stdout[7..]
  109. );
  110. } else {
  111. bail!("Invalid line obtained from remote process: '{}'", stdout);
  112. }
  113. Ok(())
  114. }
  115. pub fn check_ready(&mut self) -> Result<bool> {
  116. let mut shell_cmd = format!(
  117. "ps -p {pid} > /dev/null || (echo 'Not running'; exit 1)\n",
  118. pid = self.pid
  119. );
  120. let is_ready = match self.readiness {
  121. Readiness::IsReady => true,
  122. Readiness::WaitingForReadyFile(ref path) => {
  123. shell_cmd += &format!("rm {:?} && echo 'Ready' && exit 0\n", path);
  124. false
  125. }
  126. };
  127. shell_cmd += "echo 'Ok'";
  128. let (stdout, _stderr) = self.run_ssh(&shell_cmd)?;
  129. Ok(match stdout.trim() {
  130. "Ok" => is_ready,
  131. "Ready" => {
  132. log::info!("Remote process {} is ready", self.name);
  133. self.readiness = Readiness::IsReady;
  134. true
  135. }
  136. "Not Running" => bail!("Remote process {} is not running", self.name),
  137. output => bail!(
  138. "Unexpected output from remote process {}: {}",
  139. self.name,
  140. output
  141. ),
  142. })
  143. }
  144. pub fn kill(&mut self) -> Result<()> {
  145. let shell_cmd = match self.readiness {
  146. Readiness::IsReady => format!("pkill -P {pid}; exit 0", pid = self.pid),
  147. Readiness::WaitingForReadyFile(ref path) => format!(
  148. "pkill -P {pid}; rm -f {ready_file:?}; exit 0",
  149. pid = self.pid,
  150. ready_file = path
  151. ),
  152. };
  153. self.run_ssh(&shell_cmd)?;
  154. Ok(())
  155. }
  156. }