|
| 1 | +use std::time::Duration; |
| 2 | +use std::sync::Once; |
| 3 | +use std::process::Stdio; |
| 4 | +use std::io::{self, Read}; |
| 5 | + |
| 6 | +use rocket::fairing::AdHoc; |
| 7 | +use rocket::http::ext::IntoOwned; |
| 8 | +use rocket::http::uri::{self, Absolute, Uri}; |
| 9 | +use rocket::serde::{Deserialize, Serialize}; |
| 10 | +use rocket::{Build, Rocket}; |
| 11 | + |
| 12 | +use procspawn::SpawnError; |
| 13 | +use thiserror::Error; |
| 14 | +use ipc_channel::ipc::{IpcOneShotServer, IpcReceiver, IpcSender}; |
| 15 | + |
| 16 | +static DEFAULT_CONFIG: &str = r#" |
| 17 | + [default] |
| 18 | + address = "tcp:127.0.0.1" |
| 19 | + workers = 2 |
| 20 | + port = 0 |
| 21 | + cli_colors = false |
| 22 | + secret_key = "itlYmFR2vYKrOmFhupMIn/hyB6lYCCTXz4yaQX89XVg=" |
| 23 | +
|
| 24 | + [default.shutdown] |
| 25 | + grace = 1 |
| 26 | + mercy = 1 |
| 27 | +"#; |
| 28 | + |
| 29 | +#[derive(Debug)] |
| 30 | +#[allow(unused)] |
| 31 | +pub struct Client { |
| 32 | + client: reqwest::blocking::Client, |
| 33 | + server: procspawn::JoinHandle<()>, |
| 34 | + tls: bool, |
| 35 | + port: u16, |
| 36 | + rx: IpcReceiver<Message>, |
| 37 | +} |
| 38 | + |
| 39 | +#[derive(Error, Debug)] |
| 40 | +pub enum Error { |
| 41 | + #[error("join/kill failed: {0}")] |
| 42 | + JoinError(#[from] SpawnError), |
| 43 | + #[error("kill failed: {0}")] |
| 44 | + TermFailure(#[from] nix::errno::Errno), |
| 45 | + #[error("i/o error: {0}")] |
| 46 | + Io(#[from] io::Error), |
| 47 | + #[error("invalid URI: {0}")] |
| 48 | + Uri(#[from] uri::Error<'static>), |
| 49 | + #[error("the URI is invalid")] |
| 50 | + InvalidUri, |
| 51 | + #[error("bad request: {0}")] |
| 52 | + Request(#[from] reqwest::Error), |
| 53 | + #[error("IPC failure: {0}")] |
| 54 | + Ipc(#[from] ipc_channel::ipc::IpcError), |
| 55 | + #[error("liftoff failed")] |
| 56 | + Liftoff(String, String), |
| 57 | +} |
| 58 | + |
| 59 | +#[derive(Debug, Serialize, Deserialize)] |
| 60 | +#[serde(crate = "rocket::serde")] |
| 61 | +pub enum Message { |
| 62 | + Liftoff(bool, u16), |
| 63 | + Failure, |
| 64 | +} |
| 65 | + |
| 66 | +#[derive(Serialize, Deserialize)] |
| 67 | +#[serde(crate = "rocket::serde")] |
| 68 | +#[must_use] |
| 69 | +pub struct Token(String); |
| 70 | + |
| 71 | +pub type Result<T, E = Error> = std::result::Result<T, E>; |
| 72 | + |
| 73 | +impl Token { |
| 74 | + fn configure(&self, toml: &str, rocket: Rocket<Build>) -> Rocket<Build> { |
| 75 | + use rocket::figment::{Figment, providers::{Format, Toml}}; |
| 76 | + |
| 77 | + let toml = toml.replace("{CRATE}", env!("CARGO_MANIFEST_DIR")); |
| 78 | + let config = Figment::from(rocket.figment()) |
| 79 | + .merge(Toml::string(DEFAULT_CONFIG).nested()) |
| 80 | + .merge(Toml::string(&toml).nested()); |
| 81 | + |
| 82 | + let server = self.0.clone(); |
| 83 | + rocket.configure(config) |
| 84 | + .attach(AdHoc::on_liftoff("Liftoff", move |rocket| Box::pin(async move { |
| 85 | + let tcp = rocket.endpoints().find_map(|e| e.tcp()).unwrap(); |
| 86 | + let tls = rocket.endpoints().any(|e| e.is_tls()); |
| 87 | + let sender = IpcSender::<Message>::connect(server).unwrap(); |
| 88 | + let _ = sender.send(Message::Liftoff(tls, tcp.port())); |
| 89 | + let _ = sender.send(Message::Liftoff(tls, tcp.port())); |
| 90 | + }))) |
| 91 | + } |
| 92 | + |
| 93 | + pub fn rocket(&self, toml: &str) -> Rocket<Build> { |
| 94 | + self.configure(toml, rocket::build()) |
| 95 | + } |
| 96 | + |
| 97 | + pub fn configured_launch(self, toml: &str, rocket: Rocket<Build>) { |
| 98 | + let rocket = self.configure(toml, rocket); |
| 99 | + if let Err(e) = rocket::execute(rocket.launch()) { |
| 100 | + let sender = IpcSender::<Message>::connect(self.0).unwrap(); |
| 101 | + let _ = sender.send(Message::Failure); |
| 102 | + let _ = sender.send(Message::Failure); |
| 103 | + e.pretty_print(); |
| 104 | + std::process::exit(1); |
| 105 | + } |
| 106 | + } |
| 107 | + |
| 108 | + pub fn launch(self, rocket: Rocket<Build>) { |
| 109 | + self.configured_launch(DEFAULT_CONFIG, rocket) |
| 110 | + } |
| 111 | +} |
| 112 | +pub fn start(f: fn(Token)) -> Result<Client> { |
| 113 | + static INIT: Once = Once::new(); |
| 114 | + INIT.call_once(procspawn::init); |
| 115 | + |
| 116 | + let (ipc, server) = IpcOneShotServer::new()?; |
| 117 | + let mut server = procspawn::Builder::new() |
| 118 | + .stdin(Stdio::null()) |
| 119 | + .stdout(Stdio::piped()) |
| 120 | + .stderr(Stdio::piped()) |
| 121 | + .spawn(Token(server), f); |
| 122 | + |
| 123 | + let client = reqwest::blocking::Client::builder() |
| 124 | + .danger_accept_invalid_certs(true) |
| 125 | + .cookie_store(true) |
| 126 | + .tls_info(true) |
| 127 | + .timeout(Duration::from_secs(5)) |
| 128 | + .connect_timeout(Duration::from_secs(5)) |
| 129 | + .build()?; |
| 130 | + |
| 131 | + let (rx, _) = ipc.accept().unwrap(); |
| 132 | + match rx.recv() { |
| 133 | + Ok(Message::Liftoff(tls, port)) => Ok(Client { client, server, tls, port, rx }), |
| 134 | + Ok(Message::Failure) => { |
| 135 | + let stdout = server.stdout().unwrap(); |
| 136 | + let mut out = String::new(); |
| 137 | + stdout.read_to_string(&mut out)?; |
| 138 | + |
| 139 | + let stderr = server.stderr().unwrap(); |
| 140 | + let mut err = String::new(); |
| 141 | + stderr.read_to_string(&mut err)?; |
| 142 | + Err(Error::Liftoff(out, err)) |
| 143 | + } |
| 144 | + Err(e) => Err(e.into()), |
| 145 | + } |
| 146 | + |
| 147 | +} |
| 148 | + |
| 149 | +pub fn default() -> Result<Client> { |
| 150 | + start(|token| token.launch(rocket::build())) |
| 151 | +} |
| 152 | + |
| 153 | +impl Client { |
| 154 | + pub fn read_stdout(&mut self) -> Result<String> { |
| 155 | + let Some(stdout) = self.server.stdout() else { |
| 156 | + return Ok(String::new()); |
| 157 | + }; |
| 158 | + |
| 159 | + let mut string = String::new(); |
| 160 | + stdout.read_to_string(&mut string)?; |
| 161 | + Ok(string) |
| 162 | + } |
| 163 | + |
| 164 | + pub fn read_stderr(&mut self) -> Result<String> { |
| 165 | + let Some(stderr) = self.server.stderr() else { |
| 166 | + return Ok(String::new()); |
| 167 | + }; |
| 168 | + |
| 169 | + let mut string = String::new(); |
| 170 | + stderr.read_to_string(&mut string)?; |
| 171 | + Ok(string) |
| 172 | + } |
| 173 | + |
| 174 | + pub fn kill(&mut self) -> Result<()> { |
| 175 | + Ok(self.server.kill()?) |
| 176 | + } |
| 177 | + |
| 178 | + pub fn terminate(&mut self) -> Result<()> { |
| 179 | + use nix::{sys::signal, unistd::Pid}; |
| 180 | + |
| 181 | + let pid = Pid::from_raw(self.server.pid().unwrap() as i32); |
| 182 | + Ok(signal::kill(pid, signal::SIGTERM)?) |
| 183 | + } |
| 184 | + |
| 185 | + pub fn wait(&mut self) -> Result<()> { |
| 186 | + match self.server.join_timeout(Duration::from_secs(5)) { |
| 187 | + Ok(_) => Ok(()), |
| 188 | + Err(e) if e.is_remote_close() => Ok(()), |
| 189 | + Err(e) => Err(e.into()), |
| 190 | + } |
| 191 | + } |
| 192 | + |
| 193 | + pub fn get(&self, url: &str) -> Result<reqwest::blocking::RequestBuilder> { |
| 194 | + let uri = match Uri::parse_any(url).map_err(|e| e.into_owned())? { |
| 195 | + Uri::Origin(uri) => { |
| 196 | + let proto = if self.tls { "https" } else { "http" }; |
| 197 | + let uri = format!("{proto}://127.0.0.1:{}{uri}", self.port); |
| 198 | + Absolute::parse_owned(uri)? |
| 199 | + } |
| 200 | + Uri::Absolute(uri) => uri, |
| 201 | + _ => return Err(Error::InvalidUri), |
| 202 | + }; |
| 203 | + |
| 204 | + Ok(self.client.get(uri.to_string())) |
| 205 | + } |
| 206 | +} |
0 commit comments