##// END OF EJS Templates
rust-chg: leverage impl trait at argument position...
Yuya Nishihara -
r45183:61fda2db default
parent child Browse files
Show More
@@ -1,136 +1,128 b''
1 // Copyright 2018 Yuya Nishihara <yuya@tcha.org>
1 // Copyright 2018 Yuya Nishihara <yuya@tcha.org>
2 //
2 //
3 // This software may be used and distributed according to the terms of the
3 // This software may be used and distributed according to the terms of the
4 // GNU General Public License version 2 or any later version.
4 // GNU General Public License version 2 or any later version.
5
5
6 //! cHg extensions to command server client.
6 //! cHg extensions to command server client.
7
7
8 use bytes::{BufMut, Bytes, BytesMut};
8 use bytes::{BufMut, Bytes, BytesMut};
9 use std::ffi::OsStr;
9 use std::ffi::OsStr;
10 use std::io;
10 use std::io;
11 use std::mem;
11 use std::mem;
12 use std::os::unix::ffi::OsStrExt;
12 use std::os::unix::ffi::OsStrExt;
13 use std::os::unix::io::AsRawFd;
13 use std::os::unix::io::AsRawFd;
14 use std::path::Path;
14 use std::path::Path;
15 use tokio_hglib::protocol::{OneShotQuery, OneShotRequest};
15 use tokio_hglib::protocol::{OneShotQuery, OneShotRequest};
16 use tokio_hglib::{Client, Connection};
16 use tokio_hglib::{Client, Connection};
17
17
18 use crate::attachio::AttachIo;
18 use crate::attachio::AttachIo;
19 use crate::message::{self, Instruction};
19 use crate::message::{self, Instruction};
20 use crate::runcommand::ChgRunCommand;
20 use crate::runcommand::ChgRunCommand;
21 use crate::uihandler::SystemHandler;
21 use crate::uihandler::SystemHandler;
22
22
23 pub trait ChgClientExt<C>
23 pub trait ChgClientExt<C>
24 where
24 where
25 C: Connection + AsRawFd,
25 C: Connection + AsRawFd,
26 {
26 {
27 /// Attaches the client file descriptors to the server.
27 /// Attaches the client file descriptors to the server.
28 fn attach_io<I, O, E>(self, stdin: I, stdout: O, stderr: E) -> AttachIo<C, I, O, E>
28 fn attach_io<I, O, E>(self, stdin: I, stdout: O, stderr: E) -> AttachIo<C, I, O, E>
29 where
29 where
30 I: AsRawFd,
30 I: AsRawFd,
31 O: AsRawFd,
31 O: AsRawFd,
32 E: AsRawFd;
32 E: AsRawFd;
33
33
34 /// Changes the working directory of the server.
34 /// Changes the working directory of the server.
35 fn set_current_dir<P>(self, dir: P) -> OneShotRequest<C>
35 fn set_current_dir(self, dir: impl AsRef<Path>) -> OneShotRequest<C>;
36 where
37 P: AsRef<Path>;
38
36
39 /// Updates the environment variables of the server.
37 /// Updates the environment variables of the server.
40 fn set_env_vars_os<I, P>(self, vars: I) -> OneShotRequest<C>
38 fn set_env_vars_os(
41 where
39 self,
42 I: IntoIterator<Item = (P, P)>,
40 vars: impl IntoIterator<Item = (impl AsRef<OsStr>, impl AsRef<OsStr>)>,
43 P: AsRef<OsStr>;
41 ) -> OneShotRequest<C>;
44
42
45 /// Changes the process title of the server.
43 /// Changes the process title of the server.
46 fn set_process_name<P>(self, name: P) -> OneShotRequest<C>
44 fn set_process_name(self, name: impl AsRef<OsStr>) -> OneShotRequest<C>;
47 where
48 P: AsRef<OsStr>;
49
45
50 /// Changes the umask of the server process.
46 /// Changes the umask of the server process.
51 fn set_umask(self, mask: u32) -> OneShotRequest<C>;
47 fn set_umask(self, mask: u32) -> OneShotRequest<C>;
52
48
53 /// Runs the specified Mercurial command with cHg extension.
49 /// Runs the specified Mercurial command with cHg extension.
54 fn run_command_chg<I, P, H>(self, handler: H, args: I) -> ChgRunCommand<C, H>
50 fn run_command_chg<H>(
51 self,
52 handler: H,
53 args: impl IntoIterator<Item = impl AsRef<OsStr>>,
54 ) -> ChgRunCommand<C, H>
55 where
55 where
56 I: IntoIterator<Item = P>,
57 P: AsRef<OsStr>,
58 H: SystemHandler;
56 H: SystemHandler;
59
57
60 /// Validates if the server can run Mercurial commands with the expected
58 /// Validates if the server can run Mercurial commands with the expected
61 /// configuration.
59 /// configuration.
62 ///
60 ///
63 /// The `args` should contain early command arguments such as `--config`
61 /// The `args` should contain early command arguments such as `--config`
64 /// and `-R`.
62 /// and `-R`.
65 ///
63 ///
66 /// Client-side environment must be sent prior to this request, by
64 /// Client-side environment must be sent prior to this request, by
67 /// `set_current_dir()` and `set_env_vars_os()`.
65 /// `set_current_dir()` and `set_env_vars_os()`.
68 fn validate<I, P>(self, args: I) -> OneShotQuery<C, fn(Bytes) -> io::Result<Vec<Instruction>>>
66 fn validate(
69 where
67 self,
70 I: IntoIterator<Item = P>,
68 args: impl IntoIterator<Item = impl AsRef<OsStr>>,
71 P: AsRef<OsStr>;
69 ) -> OneShotQuery<C, fn(Bytes) -> io::Result<Vec<Instruction>>>;
72 }
70 }
73
71
74 impl<C> ChgClientExt<C> for Client<C>
72 impl<C> ChgClientExt<C> for Client<C>
75 where
73 where
76 C: Connection + AsRawFd,
74 C: Connection + AsRawFd,
77 {
75 {
78 fn attach_io<I, O, E>(self, stdin: I, stdout: O, stderr: E) -> AttachIo<C, I, O, E>
76 fn attach_io<I, O, E>(self, stdin: I, stdout: O, stderr: E) -> AttachIo<C, I, O, E>
79 where
77 where
80 I: AsRawFd,
78 I: AsRawFd,
81 O: AsRawFd,
79 O: AsRawFd,
82 E: AsRawFd,
80 E: AsRawFd,
83 {
81 {
84 AttachIo::with_client(self, stdin, stdout, Some(stderr))
82 AttachIo::with_client(self, stdin, stdout, Some(stderr))
85 }
83 }
86
84
87 fn set_current_dir<P>(self, dir: P) -> OneShotRequest<C>
85 fn set_current_dir(self, dir: impl AsRef<Path>) -> OneShotRequest<C> {
88 where
89 P: AsRef<Path>,
90 {
91 OneShotRequest::start_with_args(self, b"chdir", dir.as_ref().as_os_str().as_bytes())
86 OneShotRequest::start_with_args(self, b"chdir", dir.as_ref().as_os_str().as_bytes())
92 }
87 }
93
88
94 fn set_env_vars_os<I, P>(self, vars: I) -> OneShotRequest<C>
89 fn set_env_vars_os(
95 where
90 self,
96 I: IntoIterator<Item = (P, P)>,
91 vars: impl IntoIterator<Item = (impl AsRef<OsStr>, impl AsRef<OsStr>)>,
97 P: AsRef<OsStr>,
92 ) -> OneShotRequest<C> {
98 {
99 OneShotRequest::start_with_args(self, b"setenv", message::pack_env_vars_os(vars))
93 OneShotRequest::start_with_args(self, b"setenv", message::pack_env_vars_os(vars))
100 }
94 }
101
95
102 fn set_process_name<P>(self, name: P) -> OneShotRequest<C>
96 fn set_process_name(self, name: impl AsRef<OsStr>) -> OneShotRequest<C> {
103 where
104 P: AsRef<OsStr>,
105 {
106 OneShotRequest::start_with_args(self, b"setprocname", name.as_ref().as_bytes())
97 OneShotRequest::start_with_args(self, b"setprocname", name.as_ref().as_bytes())
107 }
98 }
108
99
109 fn set_umask(self, mask: u32) -> OneShotRequest<C> {
100 fn set_umask(self, mask: u32) -> OneShotRequest<C> {
110 let mut args = BytesMut::with_capacity(mem::size_of_val(&mask));
101 let mut args = BytesMut::with_capacity(mem::size_of_val(&mask));
111 args.put_u32_be(mask);
102 args.put_u32_be(mask);
112 OneShotRequest::start_with_args(self, b"setumask2", args)
103 OneShotRequest::start_with_args(self, b"setumask2", args)
113 }
104 }
114
105
115 fn run_command_chg<I, P, H>(self, handler: H, args: I) -> ChgRunCommand<C, H>
106 fn run_command_chg<H>(
107 self,
108 handler: H,
109 args: impl IntoIterator<Item = impl AsRef<OsStr>>,
110 ) -> ChgRunCommand<C, H>
116 where
111 where
117 I: IntoIterator<Item = P>,
118 P: AsRef<OsStr>,
119 H: SystemHandler,
112 H: SystemHandler,
120 {
113 {
121 ChgRunCommand::with_client(self, handler, message::pack_args_os(args))
114 ChgRunCommand::with_client(self, handler, message::pack_args_os(args))
122 }
115 }
123
116
124 fn validate<I, P>(self, args: I) -> OneShotQuery<C, fn(Bytes) -> io::Result<Vec<Instruction>>>
117 fn validate(
125 where
118 self,
126 I: IntoIterator<Item = P>,
119 args: impl IntoIterator<Item = impl AsRef<OsStr>>,
127 P: AsRef<OsStr>,
120 ) -> OneShotQuery<C, fn(Bytes) -> io::Result<Vec<Instruction>>> {
128 {
129 OneShotQuery::start_with_args(
121 OneShotQuery::start_with_args(
130 self,
122 self,
131 b"validate",
123 b"validate",
132 message::pack_args_os(args),
124 message::pack_args_os(args),
133 message::parse_instructions,
125 message::parse_instructions,
134 )
126 )
135 }
127 }
136 }
128 }
@@ -1,501 +1,490 b''
1 // Copyright 2011, 2018 Yuya Nishihara <yuya@tcha.org>
1 // Copyright 2011, 2018 Yuya Nishihara <yuya@tcha.org>
2 //
2 //
3 // This software may be used and distributed according to the terms of the
3 // This software may be used and distributed according to the terms of the
4 // GNU General Public License version 2 or any later version.
4 // GNU General Public License version 2 or any later version.
5
5
6 //! Utility for locating command-server process.
6 //! Utility for locating command-server process.
7
7
8 use futures::future::{self, Either, Loop};
8 use futures::future::{self, Either, Loop};
9 use log::debug;
9 use log::debug;
10 use std::env;
10 use std::env;
11 use std::ffi::{OsStr, OsString};
11 use std::ffi::{OsStr, OsString};
12 use std::fs::{self, DirBuilder};
12 use std::fs::{self, DirBuilder};
13 use std::io;
13 use std::io;
14 use std::os::unix::ffi::{OsStrExt, OsStringExt};
14 use std::os::unix::ffi::{OsStrExt, OsStringExt};
15 use std::os::unix::fs::{DirBuilderExt, MetadataExt};
15 use std::os::unix::fs::{DirBuilderExt, MetadataExt};
16 use std::path::{Path, PathBuf};
16 use std::path::{Path, PathBuf};
17 use std::process::{self, Command};
17 use std::process::{self, Command};
18 use std::time::Duration;
18 use std::time::Duration;
19 use tokio::prelude::*;
19 use tokio::prelude::*;
20 use tokio_hglib::UnixClient;
20 use tokio_hglib::UnixClient;
21 use tokio_process::{Child, CommandExt};
21 use tokio_process::{Child, CommandExt};
22 use tokio_timer;
22 use tokio_timer;
23
23
24 use crate::clientext::ChgClientExt;
24 use crate::clientext::ChgClientExt;
25 use crate::message::{Instruction, ServerSpec};
25 use crate::message::{Instruction, ServerSpec};
26 use crate::procutil;
26 use crate::procutil;
27
27
28 const REQUIRED_SERVER_CAPABILITIES: &[&str] = &[
28 const REQUIRED_SERVER_CAPABILITIES: &[&str] = &[
29 "attachio",
29 "attachio",
30 "chdir",
30 "chdir",
31 "runcommand",
31 "runcommand",
32 "setenv",
32 "setenv",
33 "setumask2",
33 "setumask2",
34 "validate",
34 "validate",
35 ];
35 ];
36
36
37 /// Helper to connect to and spawn a server process.
37 /// Helper to connect to and spawn a server process.
38 #[derive(Clone, Debug)]
38 #[derive(Clone, Debug)]
39 pub struct Locator {
39 pub struct Locator {
40 hg_command: OsString,
40 hg_command: OsString,
41 hg_early_args: Vec<OsString>,
41 hg_early_args: Vec<OsString>,
42 current_dir: PathBuf,
42 current_dir: PathBuf,
43 env_vars: Vec<(OsString, OsString)>,
43 env_vars: Vec<(OsString, OsString)>,
44 process_id: u32,
44 process_id: u32,
45 base_sock_path: PathBuf,
45 base_sock_path: PathBuf,
46 redirect_sock_path: Option<PathBuf>,
46 redirect_sock_path: Option<PathBuf>,
47 timeout: Duration,
47 timeout: Duration,
48 }
48 }
49
49
50 impl Locator {
50 impl Locator {
51 /// Creates locator capturing the current process environment.
51 /// Creates locator capturing the current process environment.
52 ///
52 ///
53 /// If no `$CHGSOCKNAME` is specified, the socket directory will be
53 /// If no `$CHGSOCKNAME` is specified, the socket directory will be
54 /// created as necessary.
54 /// created as necessary.
55 pub fn prepare_from_env() -> io::Result<Locator> {
55 pub fn prepare_from_env() -> io::Result<Locator> {
56 Ok(Locator {
56 Ok(Locator {
57 hg_command: default_hg_command(),
57 hg_command: default_hg_command(),
58 hg_early_args: Vec::new(),
58 hg_early_args: Vec::new(),
59 current_dir: env::current_dir()?,
59 current_dir: env::current_dir()?,
60 env_vars: env::vars_os().collect(),
60 env_vars: env::vars_os().collect(),
61 process_id: process::id(),
61 process_id: process::id(),
62 base_sock_path: prepare_server_socket_path()?,
62 base_sock_path: prepare_server_socket_path()?,
63 redirect_sock_path: None,
63 redirect_sock_path: None,
64 timeout: default_timeout(),
64 timeout: default_timeout(),
65 })
65 })
66 }
66 }
67
67
68 /// Temporary socket path for this client process.
68 /// Temporary socket path for this client process.
69 fn temp_sock_path(&self) -> PathBuf {
69 fn temp_sock_path(&self) -> PathBuf {
70 let src = self.base_sock_path.as_os_str().as_bytes();
70 let src = self.base_sock_path.as_os_str().as_bytes();
71 let mut buf = Vec::with_capacity(src.len() + 6); // "{src}.{pid}".len()
71 let mut buf = Vec::with_capacity(src.len() + 6); // "{src}.{pid}".len()
72 buf.extend_from_slice(src);
72 buf.extend_from_slice(src);
73 buf.extend_from_slice(format!(".{}", self.process_id).as_bytes());
73 buf.extend_from_slice(format!(".{}", self.process_id).as_bytes());
74 OsString::from_vec(buf).into()
74 OsString::from_vec(buf).into()
75 }
75 }
76
76
77 /// Specifies the arguments to be passed to the server at start.
77 /// Specifies the arguments to be passed to the server at start.
78 pub fn set_early_args<I, P>(&mut self, args: I)
78 pub fn set_early_args(&mut self, args: impl IntoIterator<Item = impl AsRef<OsStr>>) {
79 where
80 I: IntoIterator<Item = P>,
81 P: AsRef<OsStr>,
82 {
83 self.hg_early_args = args.into_iter().map(|a| a.as_ref().to_owned()).collect();
79 self.hg_early_args = args.into_iter().map(|a| a.as_ref().to_owned()).collect();
84 }
80 }
85
81
86 /// Connects to the server.
82 /// Connects to the server.
87 ///
83 ///
88 /// The server process will be spawned if not running.
84 /// The server process will be spawned if not running.
89 pub fn connect(self) -> impl Future<Item = (Self, UnixClient), Error = io::Error> {
85 pub fn connect(self) -> impl Future<Item = (Self, UnixClient), Error = io::Error> {
90 future::loop_fn((self, 0), |(loc, cnt)| {
86 future::loop_fn((self, 0), |(loc, cnt)| {
91 if cnt < 10 {
87 if cnt < 10 {
92 let fut = loc
88 let fut = loc
93 .try_connect()
89 .try_connect()
94 .and_then(|(loc, client)| {
90 .and_then(|(loc, client)| {
95 client
91 client
96 .validate(&loc.hg_early_args)
92 .validate(&loc.hg_early_args)
97 .map(|(client, instructions)| (loc, client, instructions))
93 .map(|(client, instructions)| (loc, client, instructions))
98 })
94 })
99 .and_then(move |(loc, client, instructions)| {
95 .and_then(move |(loc, client, instructions)| {
100 loc.run_instructions(client, instructions, cnt)
96 loc.run_instructions(client, instructions, cnt)
101 });
97 });
102 Either::A(fut)
98 Either::A(fut)
103 } else {
99 } else {
104 let msg = format!(
100 let msg = format!(
105 concat!(
101 concat!(
106 "too many redirections.\n",
102 "too many redirections.\n",
107 "Please make sure {:?} is not a wrapper which ",
103 "Please make sure {:?} is not a wrapper which ",
108 "changes sensitive environment variables ",
104 "changes sensitive environment variables ",
109 "before executing hg. If you have to use a ",
105 "before executing hg. If you have to use a ",
110 "wrapper, wrap chg instead of hg.",
106 "wrapper, wrap chg instead of hg.",
111 ),
107 ),
112 loc.hg_command
108 loc.hg_command
113 );
109 );
114 Either::B(future::err(io::Error::new(io::ErrorKind::Other, msg)))
110 Either::B(future::err(io::Error::new(io::ErrorKind::Other, msg)))
115 }
111 }
116 })
112 })
117 }
113 }
118
114
119 /// Runs instructions received from the server.
115 /// Runs instructions received from the server.
120 fn run_instructions(
116 fn run_instructions(
121 mut self,
117 mut self,
122 client: UnixClient,
118 client: UnixClient,
123 instructions: Vec<Instruction>,
119 instructions: Vec<Instruction>,
124 cnt: usize,
120 cnt: usize,
125 ) -> io::Result<Loop<(Self, UnixClient), (Self, usize)>> {
121 ) -> io::Result<Loop<(Self, UnixClient), (Self, usize)>> {
126 let mut reconnect = false;
122 let mut reconnect = false;
127 for inst in instructions {
123 for inst in instructions {
128 debug!("instruction: {:?}", inst);
124 debug!("instruction: {:?}", inst);
129 match inst {
125 match inst {
130 Instruction::Exit(_) => {
126 Instruction::Exit(_) => {
131 // Just returns the current connection to run the
127 // Just returns the current connection to run the
132 // unparsable command and report the error
128 // unparsable command and report the error
133 return Ok(Loop::Break((self, client)));
129 return Ok(Loop::Break((self, client)));
134 }
130 }
135 Instruction::Reconnect => {
131 Instruction::Reconnect => {
136 reconnect = true;
132 reconnect = true;
137 }
133 }
138 Instruction::Redirect(path) => {
134 Instruction::Redirect(path) => {
139 if path.parent() != self.base_sock_path.parent() {
135 if path.parent() != self.base_sock_path.parent() {
140 let msg = format!(
136 let msg = format!(
141 "insecure redirect instruction from server: {}",
137 "insecure redirect instruction from server: {}",
142 path.display()
138 path.display()
143 );
139 );
144 return Err(io::Error::new(io::ErrorKind::InvalidData, msg));
140 return Err(io::Error::new(io::ErrorKind::InvalidData, msg));
145 }
141 }
146 self.redirect_sock_path = Some(path);
142 self.redirect_sock_path = Some(path);
147 reconnect = true;
143 reconnect = true;
148 }
144 }
149 Instruction::Unlink(path) => {
145 Instruction::Unlink(path) => {
150 if path.parent() != self.base_sock_path.parent() {
146 if path.parent() != self.base_sock_path.parent() {
151 let msg = format!(
147 let msg = format!(
152 "insecure unlink instruction from server: {}",
148 "insecure unlink instruction from server: {}",
153 path.display()
149 path.display()
154 );
150 );
155 return Err(io::Error::new(io::ErrorKind::InvalidData, msg));
151 return Err(io::Error::new(io::ErrorKind::InvalidData, msg));
156 }
152 }
157 fs::remove_file(path).unwrap_or(()); // may race
153 fs::remove_file(path).unwrap_or(()); // may race
158 }
154 }
159 }
155 }
160 }
156 }
161
157
162 if reconnect {
158 if reconnect {
163 Ok(Loop::Continue((self, cnt + 1)))
159 Ok(Loop::Continue((self, cnt + 1)))
164 } else {
160 } else {
165 Ok(Loop::Break((self, client)))
161 Ok(Loop::Break((self, client)))
166 }
162 }
167 }
163 }
168
164
169 /// Tries to connect to the existing server, or spawns new if not running.
165 /// Tries to connect to the existing server, or spawns new if not running.
170 fn try_connect(self) -> impl Future<Item = (Self, UnixClient), Error = io::Error> {
166 fn try_connect(self) -> impl Future<Item = (Self, UnixClient), Error = io::Error> {
171 let sock_path = self
167 let sock_path = self
172 .redirect_sock_path
168 .redirect_sock_path
173 .as_ref()
169 .as_ref()
174 .unwrap_or(&self.base_sock_path)
170 .unwrap_or(&self.base_sock_path)
175 .clone();
171 .clone();
176 debug!("try connect to {}", sock_path.display());
172 debug!("try connect to {}", sock_path.display());
177 UnixClient::connect(sock_path)
173 UnixClient::connect(sock_path)
178 .then(|res| {
174 .then(|res| {
179 match res {
175 match res {
180 Ok(client) => Either::A(future::ok((self, client))),
176 Ok(client) => Either::A(future::ok((self, client))),
181 Err(_) => {
177 Err(_) => {
182 // Prevent us from being re-connected to the outdated
178 // Prevent us from being re-connected to the outdated
183 // master server: We were told by the server to redirect
179 // master server: We were told by the server to redirect
184 // to redirect_sock_path, which didn't work. We do not
180 // to redirect_sock_path, which didn't work. We do not
185 // want to connect to the same master server again
181 // want to connect to the same master server again
186 // because it would probably tell us the same thing.
182 // because it would probably tell us the same thing.
187 if self.redirect_sock_path.is_some() {
183 if self.redirect_sock_path.is_some() {
188 fs::remove_file(&self.base_sock_path).unwrap_or(());
184 fs::remove_file(&self.base_sock_path).unwrap_or(());
189 // may race
185 // may race
190 }
186 }
191 Either::B(self.spawn_connect())
187 Either::B(self.spawn_connect())
192 }
188 }
193 }
189 }
194 })
190 })
195 .and_then(|(loc, client)| {
191 .and_then(|(loc, client)| {
196 check_server_capabilities(client.server_spec())?;
192 check_server_capabilities(client.server_spec())?;
197 Ok((loc, client))
193 Ok((loc, client))
198 })
194 })
199 .and_then(|(loc, client)| {
195 .and_then(|(loc, client)| {
200 // It's purely optional, and the server might not support this command.
196 // It's purely optional, and the server might not support this command.
201 if client.server_spec().capabilities.contains("setprocname") {
197 if client.server_spec().capabilities.contains("setprocname") {
202 let fut = client
198 let fut = client
203 .set_process_name(format!("chg[worker/{}]", loc.process_id))
199 .set_process_name(format!("chg[worker/{}]", loc.process_id))
204 .map(|client| (loc, client));
200 .map(|client| (loc, client));
205 Either::A(fut)
201 Either::A(fut)
206 } else {
202 } else {
207 Either::B(future::ok((loc, client)))
203 Either::B(future::ok((loc, client)))
208 }
204 }
209 })
205 })
210 .and_then(|(loc, client)| {
206 .and_then(|(loc, client)| {
211 client
207 client
212 .set_current_dir(&loc.current_dir)
208 .set_current_dir(&loc.current_dir)
213 .map(|client| (loc, client))
209 .map(|client| (loc, client))
214 })
210 })
215 .and_then(|(loc, client)| {
211 .and_then(|(loc, client)| {
216 client
212 client
217 .set_env_vars_os(loc.env_vars.iter().cloned())
213 .set_env_vars_os(loc.env_vars.iter().cloned())
218 .map(|client| (loc, client))
214 .map(|client| (loc, client))
219 })
215 })
220 }
216 }
221
217
222 /// Spawns new server process and connects to it.
218 /// Spawns new server process and connects to it.
223 ///
219 ///
224 /// The server will be spawned at the current working directory, then
220 /// The server will be spawned at the current working directory, then
225 /// chdir to "/", so that the server will load configs from the target
221 /// chdir to "/", so that the server will load configs from the target
226 /// repository.
222 /// repository.
227 fn spawn_connect(self) -> impl Future<Item = (Self, UnixClient), Error = io::Error> {
223 fn spawn_connect(self) -> impl Future<Item = (Self, UnixClient), Error = io::Error> {
228 let sock_path = self.temp_sock_path();
224 let sock_path = self.temp_sock_path();
229 debug!("start cmdserver at {}", sock_path.display());
225 debug!("start cmdserver at {}", sock_path.display());
230 Command::new(&self.hg_command)
226 Command::new(&self.hg_command)
231 .arg("serve")
227 .arg("serve")
232 .arg("--cmdserver")
228 .arg("--cmdserver")
233 .arg("chgunix")
229 .arg("chgunix")
234 .arg("--address")
230 .arg("--address")
235 .arg(&sock_path)
231 .arg(&sock_path)
236 .arg("--daemon-postexec")
232 .arg("--daemon-postexec")
237 .arg("chdir:/")
233 .arg("chdir:/")
238 .args(&self.hg_early_args)
234 .args(&self.hg_early_args)
239 .current_dir(&self.current_dir)
235 .current_dir(&self.current_dir)
240 .env_clear()
236 .env_clear()
241 .envs(self.env_vars.iter().cloned())
237 .envs(self.env_vars.iter().cloned())
242 .env("CHGINTERNALMARK", "")
238 .env("CHGINTERNALMARK", "")
243 .spawn_async()
239 .spawn_async()
244 .into_future()
240 .into_future()
245 .and_then(|server| self.connect_spawned(server, sock_path))
241 .and_then(|server| self.connect_spawned(server, sock_path))
246 .and_then(|(loc, client, sock_path)| {
242 .and_then(|(loc, client, sock_path)| {
247 debug!(
243 debug!(
248 "rename {} to {}",
244 "rename {} to {}",
249 sock_path.display(),
245 sock_path.display(),
250 loc.base_sock_path.display()
246 loc.base_sock_path.display()
251 );
247 );
252 fs::rename(&sock_path, &loc.base_sock_path)?;
248 fs::rename(&sock_path, &loc.base_sock_path)?;
253 Ok((loc, client))
249 Ok((loc, client))
254 })
250 })
255 }
251 }
256
252
257 /// Tries to connect to the just spawned server repeatedly until timeout
253 /// Tries to connect to the just spawned server repeatedly until timeout
258 /// exceeded.
254 /// exceeded.
259 fn connect_spawned(
255 fn connect_spawned(
260 self,
256 self,
261 server: Child,
257 server: Child,
262 sock_path: PathBuf,
258 sock_path: PathBuf,
263 ) -> impl Future<Item = (Self, UnixClient, PathBuf), Error = io::Error> {
259 ) -> impl Future<Item = (Self, UnixClient, PathBuf), Error = io::Error> {
264 debug!("try connect to {} repeatedly", sock_path.display());
260 debug!("try connect to {} repeatedly", sock_path.display());
265 let connect = future::loop_fn(sock_path, |sock_path| {
261 let connect = future::loop_fn(sock_path, |sock_path| {
266 UnixClient::connect(sock_path.clone()).then(|res| {
262 UnixClient::connect(sock_path.clone()).then(|res| {
267 match res {
263 match res {
268 Ok(client) => Either::A(future::ok(Loop::Break((client, sock_path)))),
264 Ok(client) => Either::A(future::ok(Loop::Break((client, sock_path)))),
269 Err(_) => {
265 Err(_) => {
270 // try again with slight delay
266 // try again with slight delay
271 let fut = tokio_timer::sleep(Duration::from_millis(10))
267 let fut = tokio_timer::sleep(Duration::from_millis(10))
272 .map(|()| Loop::Continue(sock_path))
268 .map(|()| Loop::Continue(sock_path))
273 .map_err(|err| io::Error::new(io::ErrorKind::Other, err));
269 .map_err(|err| io::Error::new(io::ErrorKind::Other, err));
274 Either::B(fut)
270 Either::B(fut)
275 }
271 }
276 }
272 }
277 })
273 })
278 });
274 });
279
275
280 // waits for either connection established or server failed to start
276 // waits for either connection established or server failed to start
281 connect
277 connect
282 .select2(server)
278 .select2(server)
283 .map_err(|res| res.split().0)
279 .map_err(|res| res.split().0)
284 .timeout(self.timeout)
280 .timeout(self.timeout)
285 .map_err(|err| {
281 .map_err(|err| {
286 err.into_inner().unwrap_or_else(|| {
282 err.into_inner().unwrap_or_else(|| {
287 io::Error::new(
283 io::Error::new(
288 io::ErrorKind::TimedOut,
284 io::ErrorKind::TimedOut,
289 "timed out while connecting to server",
285 "timed out while connecting to server",
290 )
286 )
291 })
287 })
292 })
288 })
293 .and_then(|res| {
289 .and_then(|res| {
294 match res {
290 match res {
295 Either::A(((client, sock_path), server)) => {
291 Either::A(((client, sock_path), server)) => {
296 server.forget(); // continue to run in background
292 server.forget(); // continue to run in background
297 Ok((self, client, sock_path))
293 Ok((self, client, sock_path))
298 }
294 }
299 Either::B((st, _)) => Err(io::Error::new(
295 Either::B((st, _)) => Err(io::Error::new(
300 io::ErrorKind::Other,
296 io::ErrorKind::Other,
301 format!("server exited too early: {}", st),
297 format!("server exited too early: {}", st),
302 )),
298 )),
303 }
299 }
304 })
300 })
305 }
301 }
306 }
302 }
307
303
308 /// Determines the server socket to connect to.
304 /// Determines the server socket to connect to.
309 ///
305 ///
310 /// If no `$CHGSOCKNAME` is specified, the socket directory will be created
306 /// If no `$CHGSOCKNAME` is specified, the socket directory will be created
311 /// as necessary.
307 /// as necessary.
312 fn prepare_server_socket_path() -> io::Result<PathBuf> {
308 fn prepare_server_socket_path() -> io::Result<PathBuf> {
313 if let Some(s) = env::var_os("CHGSOCKNAME") {
309 if let Some(s) = env::var_os("CHGSOCKNAME") {
314 Ok(PathBuf::from(s))
310 Ok(PathBuf::from(s))
315 } else {
311 } else {
316 let mut path = default_server_socket_dir();
312 let mut path = default_server_socket_dir();
317 create_secure_dir(&path)?;
313 create_secure_dir(&path)?;
318 path.push("server");
314 path.push("server");
319 Ok(path)
315 Ok(path)
320 }
316 }
321 }
317 }
322
318
323 /// Determines the default server socket path as follows.
319 /// Determines the default server socket path as follows.
324 ///
320 ///
325 /// 1. `$XDG_RUNTIME_DIR/chg`
321 /// 1. `$XDG_RUNTIME_DIR/chg`
326 /// 2. `$TMPDIR/chg$UID`
322 /// 2. `$TMPDIR/chg$UID`
327 /// 3. `/tmp/chg$UID`
323 /// 3. `/tmp/chg$UID`
328 pub fn default_server_socket_dir() -> PathBuf {
324 pub fn default_server_socket_dir() -> PathBuf {
329 // XDG_RUNTIME_DIR should be ignored if it has an insufficient permission.
325 // XDG_RUNTIME_DIR should be ignored if it has an insufficient permission.
330 // https://standards.freedesktop.org/basedir-spec/basedir-spec-latest.html
326 // https://standards.freedesktop.org/basedir-spec/basedir-spec-latest.html
331 if let Some(Ok(s)) = env::var_os("XDG_RUNTIME_DIR").map(check_secure_dir) {
327 if let Some(Ok(s)) = env::var_os("XDG_RUNTIME_DIR").map(check_secure_dir) {
332 let mut path = PathBuf::from(s);
328 let mut path = PathBuf::from(s);
333 path.push("chg");
329 path.push("chg");
334 path
330 path
335 } else {
331 } else {
336 let mut path = env::temp_dir();
332 let mut path = env::temp_dir();
337 path.push(format!("chg{}", procutil::get_effective_uid()));
333 path.push(format!("chg{}", procutil::get_effective_uid()));
338 path
334 path
339 }
335 }
340 }
336 }
341
337
342 /// Determines the default hg command.
338 /// Determines the default hg command.
343 pub fn default_hg_command() -> OsString {
339 pub fn default_hg_command() -> OsString {
344 // TODO: maybe allow embedding the path at compile time (or load from hgrc)
340 // TODO: maybe allow embedding the path at compile time (or load from hgrc)
345 env::var_os("CHGHG")
341 env::var_os("CHGHG")
346 .or(env::var_os("HG"))
342 .or(env::var_os("HG"))
347 .unwrap_or(OsStr::new("hg").to_owned())
343 .unwrap_or(OsStr::new("hg").to_owned())
348 }
344 }
349
345
350 fn default_timeout() -> Duration {
346 fn default_timeout() -> Duration {
351 let secs = env::var("CHGTIMEOUT")
347 let secs = env::var("CHGTIMEOUT")
352 .ok()
348 .ok()
353 .and_then(|s| s.parse().ok())
349 .and_then(|s| s.parse().ok())
354 .unwrap_or(60);
350 .unwrap_or(60);
355 Duration::from_secs(secs)
351 Duration::from_secs(secs)
356 }
352 }
357
353
358 /// Creates a directory which the other users cannot access to.
354 /// Creates a directory which the other users cannot access to.
359 ///
355 ///
360 /// If the directory already exists, tests its permission.
356 /// If the directory already exists, tests its permission.
361 fn create_secure_dir<P>(path: P) -> io::Result<()>
357 fn create_secure_dir(path: impl AsRef<Path>) -> io::Result<()> {
362 where
363 P: AsRef<Path>,
364 {
365 DirBuilder::new()
358 DirBuilder::new()
366 .mode(0o700)
359 .mode(0o700)
367 .create(path.as_ref())
360 .create(path.as_ref())
368 .or_else(|err| {
361 .or_else(|err| {
369 if err.kind() == io::ErrorKind::AlreadyExists {
362 if err.kind() == io::ErrorKind::AlreadyExists {
370 check_secure_dir(path).map(|_| ())
363 check_secure_dir(path).map(|_| ())
371 } else {
364 } else {
372 Err(err)
365 Err(err)
373 }
366 }
374 })
367 })
375 }
368 }
376
369
377 fn check_secure_dir<P>(path: P) -> io::Result<P>
370 fn check_secure_dir<P>(path: P) -> io::Result<P>
378 where
371 where
379 P: AsRef<Path>,
372 P: AsRef<Path>,
380 {
373 {
381 let a = fs::symlink_metadata(path.as_ref())?;
374 let a = fs::symlink_metadata(path.as_ref())?;
382 if a.is_dir() && a.uid() == procutil::get_effective_uid() && (a.mode() & 0o777) == 0o700 {
375 if a.is_dir() && a.uid() == procutil::get_effective_uid() && (a.mode() & 0o777) == 0o700 {
383 Ok(path)
376 Ok(path)
384 } else {
377 } else {
385 Err(io::Error::new(io::ErrorKind::Other, "insecure directory"))
378 Err(io::Error::new(io::ErrorKind::Other, "insecure directory"))
386 }
379 }
387 }
380 }
388
381
389 fn check_server_capabilities(spec: &ServerSpec) -> io::Result<()> {
382 fn check_server_capabilities(spec: &ServerSpec) -> io::Result<()> {
390 let unsupported: Vec<_> = REQUIRED_SERVER_CAPABILITIES
383 let unsupported: Vec<_> = REQUIRED_SERVER_CAPABILITIES
391 .iter()
384 .iter()
392 .cloned()
385 .cloned()
393 .filter(|&s| !spec.capabilities.contains(s))
386 .filter(|&s| !spec.capabilities.contains(s))
394 .collect();
387 .collect();
395 if unsupported.is_empty() {
388 if unsupported.is_empty() {
396 Ok(())
389 Ok(())
397 } else {
390 } else {
398 let msg = format!(
391 let msg = format!(
399 "insufficient server capabilities: {}",
392 "insufficient server capabilities: {}",
400 unsupported.join(", ")
393 unsupported.join(", ")
401 );
394 );
402 Err(io::Error::new(io::ErrorKind::Other, msg))
395 Err(io::Error::new(io::ErrorKind::Other, msg))
403 }
396 }
404 }
397 }
405
398
406 /// Collects arguments which need to be passed to the server at start.
399 /// Collects arguments which need to be passed to the server at start.
407 pub fn collect_early_args<I, P>(args: I) -> Vec<OsString>
400 pub fn collect_early_args(args: impl IntoIterator<Item = impl AsRef<OsStr>>) -> Vec<OsString> {
408 where
409 I: IntoIterator<Item = P>,
410 P: AsRef<OsStr>,
411 {
412 let mut args_iter = args.into_iter();
401 let mut args_iter = args.into_iter();
413 let mut early_args = Vec::new();
402 let mut early_args = Vec::new();
414 while let Some(arg) = args_iter.next() {
403 while let Some(arg) = args_iter.next() {
415 let argb = arg.as_ref().as_bytes();
404 let argb = arg.as_ref().as_bytes();
416 if argb == b"--" {
405 if argb == b"--" {
417 break;
406 break;
418 } else if argb.starts_with(b"--") {
407 } else if argb.starts_with(b"--") {
419 let mut split = argb[2..].splitn(2, |&c| c == b'=');
408 let mut split = argb[2..].splitn(2, |&c| c == b'=');
420 match split.next().unwrap() {
409 match split.next().unwrap() {
421 b"traceback" => {
410 b"traceback" => {
422 if split.next().is_none() {
411 if split.next().is_none() {
423 early_args.push(arg.as_ref().to_owned());
412 early_args.push(arg.as_ref().to_owned());
424 }
413 }
425 }
414 }
426 b"config" | b"cwd" | b"repo" | b"repository" => {
415 b"config" | b"cwd" | b"repo" | b"repository" => {
427 if split.next().is_some() {
416 if split.next().is_some() {
428 // --<flag>=<val>
417 // --<flag>=<val>
429 early_args.push(arg.as_ref().to_owned());
418 early_args.push(arg.as_ref().to_owned());
430 } else {
419 } else {
431 // --<flag> <val>
420 // --<flag> <val>
432 args_iter.next().map(|val| {
421 args_iter.next().map(|val| {
433 early_args.push(arg.as_ref().to_owned());
422 early_args.push(arg.as_ref().to_owned());
434 early_args.push(val.as_ref().to_owned());
423 early_args.push(val.as_ref().to_owned());
435 });
424 });
436 }
425 }
437 }
426 }
438 _ => {}
427 _ => {}
439 }
428 }
440 } else if argb.starts_with(b"-R") {
429 } else if argb.starts_with(b"-R") {
441 if argb.len() > 2 {
430 if argb.len() > 2 {
442 // -R<val>
431 // -R<val>
443 early_args.push(arg.as_ref().to_owned());
432 early_args.push(arg.as_ref().to_owned());
444 } else {
433 } else {
445 // -R <val>
434 // -R <val>
446 args_iter.next().map(|val| {
435 args_iter.next().map(|val| {
447 early_args.push(arg.as_ref().to_owned());
436 early_args.push(arg.as_ref().to_owned());
448 early_args.push(val.as_ref().to_owned());
437 early_args.push(val.as_ref().to_owned());
449 });
438 });
450 }
439 }
451 }
440 }
452 }
441 }
453
442
454 early_args
443 early_args
455 }
444 }
456
445
457 #[cfg(test)]
446 #[cfg(test)]
458 mod tests {
447 mod tests {
459 use super::*;
448 use super::*;
460
449
461 #[test]
450 #[test]
462 fn collect_early_args_some() {
451 fn collect_early_args_some() {
463 assert!(collect_early_args(&[] as &[&OsStr]).is_empty());
452 assert!(collect_early_args(&[] as &[&OsStr]).is_empty());
464 assert!(collect_early_args(&["log"]).is_empty());
453 assert!(collect_early_args(&["log"]).is_empty());
465 assert_eq!(
454 assert_eq!(
466 collect_early_args(&["log", "-Ra", "foo"]),
455 collect_early_args(&["log", "-Ra", "foo"]),
467 os_string_vec_from(&[b"-Ra"])
456 os_string_vec_from(&[b"-Ra"])
468 );
457 );
469 assert_eq!(
458 assert_eq!(
470 collect_early_args(&["log", "-R", "repo", "", "--traceback", "a"]),
459 collect_early_args(&["log", "-R", "repo", "", "--traceback", "a"]),
471 os_string_vec_from(&[b"-R", b"repo", b"--traceback"])
460 os_string_vec_from(&[b"-R", b"repo", b"--traceback"])
472 );
461 );
473 assert_eq!(
462 assert_eq!(
474 collect_early_args(&["log", "--config", "diff.git=1", "-q"]),
463 collect_early_args(&["log", "--config", "diff.git=1", "-q"]),
475 os_string_vec_from(&[b"--config", b"diff.git=1"])
464 os_string_vec_from(&[b"--config", b"diff.git=1"])
476 );
465 );
477 assert_eq!(
466 assert_eq!(
478 collect_early_args(&["--cwd=..", "--repository", "r", "log"]),
467 collect_early_args(&["--cwd=..", "--repository", "r", "log"]),
479 os_string_vec_from(&[b"--cwd=..", b"--repository", b"r"])
468 os_string_vec_from(&[b"--cwd=..", b"--repository", b"r"])
480 );
469 );
481 assert_eq!(
470 assert_eq!(
482 collect_early_args(&["log", "--repo=r", "--repos", "a"]),
471 collect_early_args(&["log", "--repo=r", "--repos", "a"]),
483 os_string_vec_from(&[b"--repo=r"])
472 os_string_vec_from(&[b"--repo=r"])
484 );
473 );
485 }
474 }
486
475
487 #[test]
476 #[test]
488 fn collect_early_args_orphaned() {
477 fn collect_early_args_orphaned() {
489 assert!(collect_early_args(&["log", "-R"]).is_empty());
478 assert!(collect_early_args(&["log", "-R"]).is_empty());
490 assert!(collect_early_args(&["log", "--config"]).is_empty());
479 assert!(collect_early_args(&["log", "--config"]).is_empty());
491 }
480 }
492
481
493 #[test]
482 #[test]
494 fn collect_early_args_unwanted_value() {
483 fn collect_early_args_unwanted_value() {
495 assert!(collect_early_args(&["log", "--traceback="]).is_empty());
484 assert!(collect_early_args(&["log", "--traceback="]).is_empty());
496 }
485 }
497
486
498 fn os_string_vec_from(v: &[&[u8]]) -> Vec<OsString> {
487 fn os_string_vec_from(v: &[&[u8]]) -> Vec<OsString> {
499 v.iter().map(|s| OsStr::from_bytes(s).to_owned()).collect()
488 v.iter().map(|s| OsStr::from_bytes(s).to_owned()).collect()
500 }
489 }
501 }
490 }
@@ -1,317 +1,309 b''
1 // Copyright 2018 Yuya Nishihara <yuya@tcha.org>
1 // Copyright 2018 Yuya Nishihara <yuya@tcha.org>
2 //
2 //
3 // This software may be used and distributed according to the terms of the
3 // This software may be used and distributed according to the terms of the
4 // GNU General Public License version 2 or any later version.
4 // GNU General Public License version 2 or any later version.
5
5
6 //! Utility for parsing and building command-server messages.
6 //! Utility for parsing and building command-server messages.
7
7
8 use bytes::{BufMut, Bytes, BytesMut};
8 use bytes::{BufMut, Bytes, BytesMut};
9 use std::error;
9 use std::error;
10 use std::ffi::{OsStr, OsString};
10 use std::ffi::{OsStr, OsString};
11 use std::io;
11 use std::io;
12 use std::os::unix::ffi::OsStrExt;
12 use std::os::unix::ffi::OsStrExt;
13 use std::path::PathBuf;
13 use std::path::PathBuf;
14
14
15 pub use tokio_hglib::message::*; // re-exports
15 pub use tokio_hglib::message::*; // re-exports
16
16
17 /// Shell command type requested by the server.
17 /// Shell command type requested by the server.
18 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
18 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
19 pub enum CommandType {
19 pub enum CommandType {
20 /// Pager should be spawned.
20 /// Pager should be spawned.
21 Pager,
21 Pager,
22 /// Shell command should be executed to send back the result code.
22 /// Shell command should be executed to send back the result code.
23 System,
23 System,
24 }
24 }
25
25
26 /// Shell command requested by the server.
26 /// Shell command requested by the server.
27 #[derive(Clone, Debug, Eq, PartialEq)]
27 #[derive(Clone, Debug, Eq, PartialEq)]
28 pub struct CommandSpec {
28 pub struct CommandSpec {
29 pub command: OsString,
29 pub command: OsString,
30 pub current_dir: OsString,
30 pub current_dir: OsString,
31 pub envs: Vec<(OsString, OsString)>,
31 pub envs: Vec<(OsString, OsString)>,
32 }
32 }
33
33
34 /// Parses "S" channel request into command type and spec.
34 /// Parses "S" channel request into command type and spec.
35 pub fn parse_command_spec(data: Bytes) -> io::Result<(CommandType, CommandSpec)> {
35 pub fn parse_command_spec(data: Bytes) -> io::Result<(CommandType, CommandSpec)> {
36 let mut split = data.split(|&c| c == b'\0');
36 let mut split = data.split(|&c| c == b'\0');
37 let ctype = parse_command_type(split.next().ok_or(new_parse_error("missing type"))?)?;
37 let ctype = parse_command_type(split.next().ok_or(new_parse_error("missing type"))?)?;
38 let command = split.next().ok_or(new_parse_error("missing command"))?;
38 let command = split.next().ok_or(new_parse_error("missing command"))?;
39 let current_dir = split.next().ok_or(new_parse_error("missing current dir"))?;
39 let current_dir = split.next().ok_or(new_parse_error("missing current dir"))?;
40
40
41 let mut envs = Vec::new();
41 let mut envs = Vec::new();
42 for l in split {
42 for l in split {
43 let mut s = l.splitn(2, |&c| c == b'=');
43 let mut s = l.splitn(2, |&c| c == b'=');
44 let k = s.next().unwrap();
44 let k = s.next().unwrap();
45 let v = s.next().ok_or(new_parse_error("malformed env"))?;
45 let v = s.next().ok_or(new_parse_error("malformed env"))?;
46 envs.push((
46 envs.push((
47 OsStr::from_bytes(k).to_owned(),
47 OsStr::from_bytes(k).to_owned(),
48 OsStr::from_bytes(v).to_owned(),
48 OsStr::from_bytes(v).to_owned(),
49 ));
49 ));
50 }
50 }
51
51
52 let spec = CommandSpec {
52 let spec = CommandSpec {
53 command: OsStr::from_bytes(command).to_owned(),
53 command: OsStr::from_bytes(command).to_owned(),
54 current_dir: OsStr::from_bytes(current_dir).to_owned(),
54 current_dir: OsStr::from_bytes(current_dir).to_owned(),
55 envs: envs,
55 envs: envs,
56 };
56 };
57 Ok((ctype, spec))
57 Ok((ctype, spec))
58 }
58 }
59
59
60 fn parse_command_type(value: &[u8]) -> io::Result<CommandType> {
60 fn parse_command_type(value: &[u8]) -> io::Result<CommandType> {
61 match value {
61 match value {
62 b"pager" => Ok(CommandType::Pager),
62 b"pager" => Ok(CommandType::Pager),
63 b"system" => Ok(CommandType::System),
63 b"system" => Ok(CommandType::System),
64 _ => Err(new_parse_error(format!(
64 _ => Err(new_parse_error(format!(
65 "unknown command type: {}",
65 "unknown command type: {}",
66 decode_latin1(value)
66 decode_latin1(value)
67 ))),
67 ))),
68 }
68 }
69 }
69 }
70
70
71 /// Client-side instruction requested by the server.
71 /// Client-side instruction requested by the server.
72 #[derive(Clone, Debug, Eq, PartialEq)]
72 #[derive(Clone, Debug, Eq, PartialEq)]
73 pub enum Instruction {
73 pub enum Instruction {
74 Exit(i32),
74 Exit(i32),
75 Reconnect,
75 Reconnect,
76 Redirect(PathBuf),
76 Redirect(PathBuf),
77 Unlink(PathBuf),
77 Unlink(PathBuf),
78 }
78 }
79
79
80 /// Parses validation result into instructions.
80 /// Parses validation result into instructions.
81 pub fn parse_instructions(data: Bytes) -> io::Result<Vec<Instruction>> {
81 pub fn parse_instructions(data: Bytes) -> io::Result<Vec<Instruction>> {
82 let mut instructions = Vec::new();
82 let mut instructions = Vec::new();
83 for l in data.split(|&c| c == b'\0') {
83 for l in data.split(|&c| c == b'\0') {
84 if l.is_empty() {
84 if l.is_empty() {
85 continue;
85 continue;
86 }
86 }
87 let mut s = l.splitn(2, |&c| c == b' ');
87 let mut s = l.splitn(2, |&c| c == b' ');
88 let inst = match (s.next().unwrap(), s.next()) {
88 let inst = match (s.next().unwrap(), s.next()) {
89 (b"exit", Some(arg)) => decode_latin1(arg)
89 (b"exit", Some(arg)) => decode_latin1(arg)
90 .parse()
90 .parse()
91 .map(Instruction::Exit)
91 .map(Instruction::Exit)
92 .map_err(|_| new_parse_error(format!("invalid exit code: {:?}", arg)))?,
92 .map_err(|_| new_parse_error(format!("invalid exit code: {:?}", arg)))?,
93 (b"reconnect", None) => Instruction::Reconnect,
93 (b"reconnect", None) => Instruction::Reconnect,
94 (b"redirect", Some(arg)) => {
94 (b"redirect", Some(arg)) => {
95 Instruction::Redirect(OsStr::from_bytes(arg).to_owned().into())
95 Instruction::Redirect(OsStr::from_bytes(arg).to_owned().into())
96 }
96 }
97 (b"unlink", Some(arg)) => Instruction::Unlink(OsStr::from_bytes(arg).to_owned().into()),
97 (b"unlink", Some(arg)) => Instruction::Unlink(OsStr::from_bytes(arg).to_owned().into()),
98 _ => {
98 _ => {
99 return Err(new_parse_error(format!("unknown command: {:?}", l)));
99 return Err(new_parse_error(format!("unknown command: {:?}", l)));
100 }
100 }
101 };
101 };
102 instructions.push(inst);
102 instructions.push(inst);
103 }
103 }
104 Ok(instructions)
104 Ok(instructions)
105 }
105 }
106
106
107 // allocate large buffer as environment variables can be quite long
107 // allocate large buffer as environment variables can be quite long
108 const INITIAL_PACKED_ENV_VARS_CAPACITY: usize = 4096;
108 const INITIAL_PACKED_ENV_VARS_CAPACITY: usize = 4096;
109
109
110 /// Packs environment variables of platform encoding into bytes.
110 /// Packs environment variables of platform encoding into bytes.
111 ///
111 ///
112 /// # Panics
112 /// # Panics
113 ///
113 ///
114 /// Panics if key or value contains `\0` character, or key contains '='
114 /// Panics if key or value contains `\0` character, or key contains '='
115 /// character.
115 /// character.
116 pub fn pack_env_vars_os<I, P>(vars: I) -> Bytes
116 pub fn pack_env_vars_os(
117 where
117 vars: impl IntoIterator<Item = (impl AsRef<OsStr>, impl AsRef<OsStr>)>,
118 I: IntoIterator<Item = (P, P)>,
118 ) -> Bytes {
119 P: AsRef<OsStr>,
120 {
121 let mut vars_iter = vars.into_iter();
119 let mut vars_iter = vars.into_iter();
122 if let Some((k, v)) = vars_iter.next() {
120 if let Some((k, v)) = vars_iter.next() {
123 let mut dst = BytesMut::with_capacity(INITIAL_PACKED_ENV_VARS_CAPACITY);
121 let mut dst = BytesMut::with_capacity(INITIAL_PACKED_ENV_VARS_CAPACITY);
124 pack_env_into(&mut dst, k.as_ref(), v.as_ref());
122 pack_env_into(&mut dst, k.as_ref(), v.as_ref());
125 for (k, v) in vars_iter {
123 for (k, v) in vars_iter {
126 dst.reserve(1);
124 dst.reserve(1);
127 dst.put_u8(b'\0');
125 dst.put_u8(b'\0');
128 pack_env_into(&mut dst, k.as_ref(), v.as_ref());
126 pack_env_into(&mut dst, k.as_ref(), v.as_ref());
129 }
127 }
130 dst.freeze()
128 dst.freeze()
131 } else {
129 } else {
132 Bytes::new()
130 Bytes::new()
133 }
131 }
134 }
132 }
135
133
136 fn pack_env_into(dst: &mut BytesMut, k: &OsStr, v: &OsStr) {
134 fn pack_env_into(dst: &mut BytesMut, k: &OsStr, v: &OsStr) {
137 assert!(!k.as_bytes().contains(&0), "key shouldn't contain NUL");
135 assert!(!k.as_bytes().contains(&0), "key shouldn't contain NUL");
138 assert!(!k.as_bytes().contains(&b'='), "key shouldn't contain '='");
136 assert!(!k.as_bytes().contains(&b'='), "key shouldn't contain '='");
139 assert!(!v.as_bytes().contains(&0), "value shouldn't contain NUL");
137 assert!(!v.as_bytes().contains(&0), "value shouldn't contain NUL");
140 dst.reserve(k.as_bytes().len() + 1 + v.as_bytes().len());
138 dst.reserve(k.as_bytes().len() + 1 + v.as_bytes().len());
141 dst.put_slice(k.as_bytes());
139 dst.put_slice(k.as_bytes());
142 dst.put_u8(b'=');
140 dst.put_u8(b'=');
143 dst.put_slice(v.as_bytes());
141 dst.put_slice(v.as_bytes());
144 }
142 }
145
143
146 fn decode_latin1<S>(s: S) -> String
144 fn decode_latin1(s: impl AsRef<[u8]>) -> String {
147 where
148 S: AsRef<[u8]>,
149 {
150 s.as_ref().iter().map(|&c| c as char).collect()
145 s.as_ref().iter().map(|&c| c as char).collect()
151 }
146 }
152
147
153 fn new_parse_error<E>(error: E) -> io::Error
148 fn new_parse_error(error: impl Into<Box<dyn error::Error + Send + Sync>>) -> io::Error {
154 where
155 E: Into<Box<dyn error::Error + Send + Sync>>,
156 {
157 io::Error::new(io::ErrorKind::InvalidData, error)
149 io::Error::new(io::ErrorKind::InvalidData, error)
158 }
150 }
159
151
160 #[cfg(test)]
152 #[cfg(test)]
161 mod tests {
153 mod tests {
162 use super::*;
154 use super::*;
163 use std::os::unix::ffi::OsStringExt;
155 use std::os::unix::ffi::OsStringExt;
164 use std::panic;
156 use std::panic;
165
157
166 #[test]
158 #[test]
167 fn parse_command_spec_good() {
159 fn parse_command_spec_good() {
168 let src = [
160 let src = [
169 b"pager".as_ref(),
161 b"pager".as_ref(),
170 b"less -FRX".as_ref(),
162 b"less -FRX".as_ref(),
171 b"/tmp".as_ref(),
163 b"/tmp".as_ref(),
172 b"LANG=C".as_ref(),
164 b"LANG=C".as_ref(),
173 b"HGPLAIN=".as_ref(),
165 b"HGPLAIN=".as_ref(),
174 ]
166 ]
175 .join(&0);
167 .join(&0);
176 let spec = CommandSpec {
168 let spec = CommandSpec {
177 command: os_string_from(b"less -FRX"),
169 command: os_string_from(b"less -FRX"),
178 current_dir: os_string_from(b"/tmp"),
170 current_dir: os_string_from(b"/tmp"),
179 envs: vec![
171 envs: vec![
180 (os_string_from(b"LANG"), os_string_from(b"C")),
172 (os_string_from(b"LANG"), os_string_from(b"C")),
181 (os_string_from(b"HGPLAIN"), os_string_from(b"")),
173 (os_string_from(b"HGPLAIN"), os_string_from(b"")),
182 ],
174 ],
183 };
175 };
184 assert_eq!(
176 assert_eq!(
185 parse_command_spec(Bytes::from(src)).unwrap(),
177 parse_command_spec(Bytes::from(src)).unwrap(),
186 (CommandType::Pager, spec)
178 (CommandType::Pager, spec)
187 );
179 );
188 }
180 }
189
181
190 #[test]
182 #[test]
191 fn parse_command_spec_too_short() {
183 fn parse_command_spec_too_short() {
192 assert!(parse_command_spec(Bytes::from_static(b"")).is_err());
184 assert!(parse_command_spec(Bytes::from_static(b"")).is_err());
193 assert!(parse_command_spec(Bytes::from_static(b"pager")).is_err());
185 assert!(parse_command_spec(Bytes::from_static(b"pager")).is_err());
194 assert!(parse_command_spec(Bytes::from_static(b"pager\0less")).is_err());
186 assert!(parse_command_spec(Bytes::from_static(b"pager\0less")).is_err());
195 }
187 }
196
188
197 #[test]
189 #[test]
198 fn parse_command_spec_malformed_env() {
190 fn parse_command_spec_malformed_env() {
199 assert!(parse_command_spec(Bytes::from_static(b"pager\0less\0/tmp\0HOME")).is_err());
191 assert!(parse_command_spec(Bytes::from_static(b"pager\0less\0/tmp\0HOME")).is_err());
200 }
192 }
201
193
202 #[test]
194 #[test]
203 fn parse_command_spec_unknown_type() {
195 fn parse_command_spec_unknown_type() {
204 assert!(parse_command_spec(Bytes::from_static(b"paper\0less")).is_err());
196 assert!(parse_command_spec(Bytes::from_static(b"paper\0less")).is_err());
205 }
197 }
206
198
207 #[test]
199 #[test]
208 fn parse_instructions_good() {
200 fn parse_instructions_good() {
209 let src = [
201 let src = [
210 b"exit 123".as_ref(),
202 b"exit 123".as_ref(),
211 b"reconnect".as_ref(),
203 b"reconnect".as_ref(),
212 b"redirect /whatever".as_ref(),
204 b"redirect /whatever".as_ref(),
213 b"unlink /someother".as_ref(),
205 b"unlink /someother".as_ref(),
214 ]
206 ]
215 .join(&0);
207 .join(&0);
216 let insts = vec![
208 let insts = vec![
217 Instruction::Exit(123),
209 Instruction::Exit(123),
218 Instruction::Reconnect,
210 Instruction::Reconnect,
219 Instruction::Redirect(path_buf_from(b"/whatever")),
211 Instruction::Redirect(path_buf_from(b"/whatever")),
220 Instruction::Unlink(path_buf_from(b"/someother")),
212 Instruction::Unlink(path_buf_from(b"/someother")),
221 ];
213 ];
222 assert_eq!(parse_instructions(Bytes::from(src)).unwrap(), insts);
214 assert_eq!(parse_instructions(Bytes::from(src)).unwrap(), insts);
223 }
215 }
224
216
225 #[test]
217 #[test]
226 fn parse_instructions_empty() {
218 fn parse_instructions_empty() {
227 assert_eq!(parse_instructions(Bytes::new()).unwrap(), vec![]);
219 assert_eq!(parse_instructions(Bytes::new()).unwrap(), vec![]);
228 assert_eq!(
220 assert_eq!(
229 parse_instructions(Bytes::from_static(b"\0")).unwrap(),
221 parse_instructions(Bytes::from_static(b"\0")).unwrap(),
230 vec![]
222 vec![]
231 );
223 );
232 }
224 }
233
225
234 #[test]
226 #[test]
235 fn parse_instructions_malformed_exit_code() {
227 fn parse_instructions_malformed_exit_code() {
236 assert!(parse_instructions(Bytes::from_static(b"exit foo")).is_err());
228 assert!(parse_instructions(Bytes::from_static(b"exit foo")).is_err());
237 }
229 }
238
230
239 #[test]
231 #[test]
240 fn parse_instructions_missing_argument() {
232 fn parse_instructions_missing_argument() {
241 assert!(parse_instructions(Bytes::from_static(b"exit")).is_err());
233 assert!(parse_instructions(Bytes::from_static(b"exit")).is_err());
242 assert!(parse_instructions(Bytes::from_static(b"redirect")).is_err());
234 assert!(parse_instructions(Bytes::from_static(b"redirect")).is_err());
243 assert!(parse_instructions(Bytes::from_static(b"unlink")).is_err());
235 assert!(parse_instructions(Bytes::from_static(b"unlink")).is_err());
244 }
236 }
245
237
246 #[test]
238 #[test]
247 fn parse_instructions_unknown_command() {
239 fn parse_instructions_unknown_command() {
248 assert!(parse_instructions(Bytes::from_static(b"quit 0")).is_err());
240 assert!(parse_instructions(Bytes::from_static(b"quit 0")).is_err());
249 }
241 }
250
242
251 #[test]
243 #[test]
252 fn pack_env_vars_os_good() {
244 fn pack_env_vars_os_good() {
253 assert_eq!(
245 assert_eq!(
254 pack_env_vars_os(vec![] as Vec<(OsString, OsString)>),
246 pack_env_vars_os(vec![] as Vec<(OsString, OsString)>),
255 Bytes::new()
247 Bytes::new()
256 );
248 );
257 assert_eq!(
249 assert_eq!(
258 pack_env_vars_os(vec![os_string_pair_from(b"FOO", b"bar")]),
250 pack_env_vars_os(vec![os_string_pair_from(b"FOO", b"bar")]),
259 Bytes::from_static(b"FOO=bar")
251 Bytes::from_static(b"FOO=bar")
260 );
252 );
261 assert_eq!(
253 assert_eq!(
262 pack_env_vars_os(vec![
254 pack_env_vars_os(vec![
263 os_string_pair_from(b"FOO", b""),
255 os_string_pair_from(b"FOO", b""),
264 os_string_pair_from(b"BAR", b"baz")
256 os_string_pair_from(b"BAR", b"baz")
265 ]),
257 ]),
266 Bytes::from_static(b"FOO=\0BAR=baz")
258 Bytes::from_static(b"FOO=\0BAR=baz")
267 );
259 );
268 }
260 }
269
261
270 #[test]
262 #[test]
271 fn pack_env_vars_os_large_key() {
263 fn pack_env_vars_os_large_key() {
272 let mut buf = vec![b'A'; INITIAL_PACKED_ENV_VARS_CAPACITY];
264 let mut buf = vec![b'A'; INITIAL_PACKED_ENV_VARS_CAPACITY];
273 let envs = vec![os_string_pair_from(&buf, b"")];
265 let envs = vec![os_string_pair_from(&buf, b"")];
274 buf.push(b'=');
266 buf.push(b'=');
275 assert_eq!(pack_env_vars_os(envs), Bytes::from(buf));
267 assert_eq!(pack_env_vars_os(envs), Bytes::from(buf));
276 }
268 }
277
269
278 #[test]
270 #[test]
279 fn pack_env_vars_os_large_value() {
271 fn pack_env_vars_os_large_value() {
280 let mut buf = vec![b'A', b'='];
272 let mut buf = vec![b'A', b'='];
281 buf.resize(INITIAL_PACKED_ENV_VARS_CAPACITY + 1, b'a');
273 buf.resize(INITIAL_PACKED_ENV_VARS_CAPACITY + 1, b'a');
282 let envs = vec![os_string_pair_from(&buf[..1], &buf[2..])];
274 let envs = vec![os_string_pair_from(&buf[..1], &buf[2..])];
283 assert_eq!(pack_env_vars_os(envs), Bytes::from(buf));
275 assert_eq!(pack_env_vars_os(envs), Bytes::from(buf));
284 }
276 }
285
277
286 #[test]
278 #[test]
287 fn pack_env_vars_os_nul_eq() {
279 fn pack_env_vars_os_nul_eq() {
288 assert!(panic::catch_unwind(|| {
280 assert!(panic::catch_unwind(|| {
289 pack_env_vars_os(vec![os_string_pair_from(b"\0", b"")])
281 pack_env_vars_os(vec![os_string_pair_from(b"\0", b"")])
290 })
282 })
291 .is_err());
283 .is_err());
292 assert!(panic::catch_unwind(|| {
284 assert!(panic::catch_unwind(|| {
293 pack_env_vars_os(vec![os_string_pair_from(b"FOO", b"\0bar")])
285 pack_env_vars_os(vec![os_string_pair_from(b"FOO", b"\0bar")])
294 })
286 })
295 .is_err());
287 .is_err());
296 assert!(panic::catch_unwind(|| {
288 assert!(panic::catch_unwind(|| {
297 pack_env_vars_os(vec![os_string_pair_from(b"FO=", b"bar")])
289 pack_env_vars_os(vec![os_string_pair_from(b"FO=", b"bar")])
298 })
290 })
299 .is_err());
291 .is_err());
300 assert_eq!(
292 assert_eq!(
301 pack_env_vars_os(vec![os_string_pair_from(b"FOO", b"=ba")]),
293 pack_env_vars_os(vec![os_string_pair_from(b"FOO", b"=ba")]),
302 Bytes::from_static(b"FOO==ba")
294 Bytes::from_static(b"FOO==ba")
303 );
295 );
304 }
296 }
305
297
306 fn os_string_from(s: &[u8]) -> OsString {
298 fn os_string_from(s: &[u8]) -> OsString {
307 OsString::from_vec(s.to_vec())
299 OsString::from_vec(s.to_vec())
308 }
300 }
309
301
310 fn os_string_pair_from(k: &[u8], v: &[u8]) -> (OsString, OsString) {
302 fn os_string_pair_from(k: &[u8], v: &[u8]) -> (OsString, OsString) {
311 (os_string_from(k), os_string_from(v))
303 (os_string_from(k), os_string_from(v))
312 }
304 }
313
305
314 fn path_buf_from(s: &[u8]) -> PathBuf {
306 fn path_buf_from(s: &[u8]) -> PathBuf {
315 os_string_from(s).into()
307 os_string_from(s).into()
316 }
308 }
317 }
309 }
General Comments 0
You need to be logged in to leave comments. Login now