]> git.proxmox.com Git - proxmox-backup.git/blob - proxmox-rest-server/src/command_socket.rs
rename CommandoSocket to CommandSocket
[proxmox-backup.git] / proxmox-rest-server / src / command_socket.rs
1 use anyhow::{bail, format_err, Error};
2
3 use std::collections::HashMap;
4 use std::os::unix::io::AsRawFd;
5 use std::path::{PathBuf, Path};
6 use std::sync::Arc;
7
8 use futures::*;
9 use tokio::net::UnixListener;
10 use serde::Serialize;
11 use serde_json::Value;
12 use nix::sys::socket;
13 use nix::unistd::Gid;
14
15 // Listens on a Unix Socket to handle simple command asynchronously
16 fn create_control_socket<P, F>(path: P, gid: Gid, func: F) -> Result<impl Future<Output = ()>, Error>
17 where
18 P: Into<PathBuf>,
19 F: Fn(Value) -> Result<Value, Error> + Send + Sync + 'static,
20 {
21 let path: PathBuf = path.into();
22
23 let gid = gid.as_raw();
24
25 let socket = UnixListener::bind(&path)?;
26
27 let func = Arc::new(func);
28
29 let control_future = async move {
30 loop {
31 let (conn, _addr) = match socket.accept().await {
32 Ok(data) => data,
33 Err(err) => {
34 eprintln!("failed to accept on control socket {:?}: {}", path, err);
35 continue;
36 }
37 };
38
39 let opt = socket::sockopt::PeerCredentials {};
40 let cred = match socket::getsockopt(conn.as_raw_fd(), opt) {
41 Ok(cred) => cred,
42 Err(err) => {
43 eprintln!("no permissions - unable to read peer credential - {}", err);
44 continue;
45 }
46 };
47
48 // check permissions (same gid, root user, or backup group)
49 let mygid = unsafe { libc::getgid() };
50 if !(cred.uid() == 0 || cred.gid() == mygid || cred.gid() == gid) {
51 eprintln!("no permissions for {:?}", cred);
52 continue;
53 }
54
55 let (rx, mut tx) = tokio::io::split(conn);
56
57 let abort_future = super::last_worker_future().map(|_| ());
58
59 use tokio::io::{AsyncBufReadExt, AsyncWriteExt};
60 let func = Arc::clone(&func);
61 let path = path.clone();
62 tokio::spawn(futures::future::select(
63 async move {
64 let mut rx = tokio::io::BufReader::new(rx);
65 let mut line = String::new();
66 loop {
67 line.clear();
68 match rx.read_line({ line.clear(); &mut line }).await {
69 Ok(0) => break,
70 Ok(_) => (),
71 Err(err) => {
72 eprintln!("control socket {:?} read error: {}", path, err);
73 return;
74 }
75 }
76
77 let response = match line.parse::<Value>() {
78 Ok(param) => match func(param) {
79 Ok(res) => format!("OK: {}\n", res),
80 Err(err) => format!("ERROR: {}\n", err),
81 }
82 Err(err) => format!("ERROR: {}\n", err),
83 };
84
85 if let Err(err) = tx.write_all(response.as_bytes()).await {
86 eprintln!("control socket {:?} write response error: {}", path, err);
87 return;
88 }
89 }
90 }.boxed(),
91 abort_future,
92 ).map(|_| ()));
93 }
94 }.boxed();
95
96 let abort_future = crate::last_worker_future().map_err(|_| {});
97 let task = futures::future::select(
98 control_future,
99 abort_future,
100 ).map(|_: futures::future::Either<(Result<(), Error>, _), _>| ());
101
102 Ok(task)
103 }
104
105 /// Send a command to the specified socket
106 pub async fn send_command<P, T>(path: P, params: &T) -> Result<Value, Error>
107 where
108 P: AsRef<Path>,
109 T: ?Sized + Serialize,
110 {
111 let mut command_string = serde_json::to_string(params)?;
112 command_string.push('\n');
113 send_raw_command(path.as_ref(), &command_string).await
114 }
115
116 /// Send a raw command (string) to the specified socket
117 pub async fn send_raw_command<P>(path: P, command_string: &str) -> Result<Value, Error>
118 where
119 P: AsRef<Path>,
120 {
121 use tokio::io::{AsyncBufReadExt, AsyncWriteExt};
122
123 let mut conn = tokio::net::UnixStream::connect(path)
124 .map_err(move |err| format_err!("control socket connect failed - {}", err))
125 .await?;
126
127 conn.write_all(command_string.as_bytes()).await?;
128 if !command_string.as_bytes().ends_with(b"\n") {
129 conn.write_all(b"\n").await?;
130 }
131
132 AsyncWriteExt::shutdown(&mut conn).await?;
133 let mut rx = tokio::io::BufReader::new(conn);
134 let mut data = String::new();
135 if rx.read_line(&mut data).await? == 0 {
136 bail!("no response");
137 }
138 if let Some(res) = data.strip_prefix("OK: ") {
139 match res.parse::<Value>() {
140 Ok(v) => Ok(v),
141 Err(err) => bail!("unable to parse json response - {}", err),
142 }
143 } else if let Some(err) = data.strip_prefix("ERROR: ") {
144 bail!("{}", err);
145 } else {
146 bail!("unable to parse response: {}", data);
147 }
148 }
149
150 // A callback for a specific commando socket.
151 type CommandSocketFn = Box<(dyn Fn(Option<&Value>) -> Result<Value, Error> + Send + Sync + 'static)>;
152
153 /// Tooling to get a single control command socket where one can
154 /// register multiple commands dynamically.
155 ///
156 /// The socket is activated by calling [spawn](CommandSocket::spawn),
157 /// which spawns an async tokio task to process the commands.
158 pub struct CommandSocket {
159 socket: PathBuf,
160 gid: Gid,
161 commands: HashMap<String, CommandSocketFn>,
162 }
163
164 impl CommandSocket {
165 /// Creates a new instance.
166 pub fn new<P>(path: P, gid: Gid) -> Self
167 where P: Into<PathBuf>,
168 {
169 CommandSocket {
170 socket: path.into(),
171 gid,
172 commands: HashMap::new(),
173 }
174 }
175
176 /// Spawn the socket and consume self, meaning you cannot register commands anymore after
177 /// calling this.
178 pub fn spawn(self) -> Result<(), Error> {
179 let control_future = create_control_socket(self.socket.to_owned(), self.gid, move |param| {
180 let param = param
181 .as_object()
182 .ok_or_else(|| format_err!("unable to parse parameters (expected json object)"))?;
183
184 let command = match param.get("command") {
185 Some(Value::String(command)) => command.as_str(),
186 None => bail!("no command"),
187 _ => bail!("unable to parse command"),
188 };
189
190 if !self.commands.contains_key(command) {
191 bail!("got unknown command '{}'", command);
192 }
193
194 match self.commands.get(command) {
195 None => bail!("got unknown command '{}'", command),
196 Some(handler) => {
197 let args = param.get("args"); //.unwrap_or(&Value::Null);
198 (handler)(args)
199 },
200 }
201 })?;
202
203 tokio::spawn(control_future);
204
205 Ok(())
206 }
207
208 /// Register a new command with a callback.
209 pub fn register_command<F>(
210 &mut self,
211 command: String,
212 handler: F,
213 ) -> Result<(), Error>
214 where
215 F: Fn(Option<&Value>) -> Result<Value, Error> + Send + Sync + 'static,
216 {
217
218 if self.commands.contains_key(&command) {
219 bail!("command '{}' already exists!", command);
220 }
221
222 self.commands.insert(command, Box::new(handler));
223
224 Ok(())
225 }
226 }