Show More
@@ -5,7 +5,7 b'' | |||||
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::Bytes; |
|
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; | |
@@ -67,6 +67,45 b' fn parse_command_type(value: &[u8]) -> i' | |||||
67 | } |
|
67 | } | |
68 | } |
|
68 | } | |
69 |
|
69 | |||
|
70 | // allocate large buffer as environment variables can be quite long | |||
|
71 | const INITIAL_PACKED_ENV_VARS_CAPACITY: usize = 4096; | |||
|
72 | ||||
|
73 | /// Packs environment variables of platform encoding into bytes. | |||
|
74 | /// | |||
|
75 | /// # Panics | |||
|
76 | /// | |||
|
77 | /// Panics if key or value contains `\0` character, or key contains '=' | |||
|
78 | /// character. | |||
|
79 | pub fn pack_env_vars_os<I, P>(vars: I) -> Bytes | |||
|
80 | where | |||
|
81 | I: IntoIterator<Item = (P, P)>, | |||
|
82 | P: AsRef<OsStr>, | |||
|
83 | { | |||
|
84 | let mut vars_iter = vars.into_iter(); | |||
|
85 | if let Some((k, v)) = vars_iter.next() { | |||
|
86 | let mut dst = BytesMut::with_capacity(INITIAL_PACKED_ENV_VARS_CAPACITY); | |||
|
87 | pack_env_into(&mut dst, k.as_ref(), v.as_ref()); | |||
|
88 | for (k, v) in vars_iter { | |||
|
89 | dst.reserve(1); | |||
|
90 | dst.put_u8(b'\0'); | |||
|
91 | pack_env_into(&mut dst, k.as_ref(), v.as_ref()); | |||
|
92 | } | |||
|
93 | dst.freeze() | |||
|
94 | } else { | |||
|
95 | Bytes::new() | |||
|
96 | } | |||
|
97 | } | |||
|
98 | ||||
|
99 | fn pack_env_into(dst: &mut BytesMut, k: &OsStr, v: &OsStr) { | |||
|
100 | assert!(!k.as_bytes().contains(&0), "key shouldn't contain NUL"); | |||
|
101 | assert!(!k.as_bytes().contains(&b'='), "key shouldn't contain '='"); | |||
|
102 | assert!(!v.as_bytes().contains(&0), "value shouldn't contain NUL"); | |||
|
103 | dst.reserve(k.as_bytes().len() + 1 + v.as_bytes().len()); | |||
|
104 | dst.put_slice(k.as_bytes()); | |||
|
105 | dst.put_u8(b'='); | |||
|
106 | dst.put_slice(v.as_bytes()); | |||
|
107 | } | |||
|
108 | ||||
70 | fn decode_latin1<S>(s: S) -> String |
|
109 | fn decode_latin1<S>(s: S) -> String | |
71 | where |
|
110 | where | |
72 | S: AsRef<[u8]>, |
|
111 | S: AsRef<[u8]>, | |
@@ -85,6 +124,7 b' where' | |||||
85 | mod tests { |
|
124 | mod tests { | |
86 | use super::*; |
|
125 | use super::*; | |
87 | use std::os::unix::ffi::OsStringExt; |
|
126 | use std::os::unix::ffi::OsStringExt; | |
|
127 | use std::panic; | |||
88 |
|
128 | |||
89 | #[test] |
|
129 | #[test] | |
90 | fn parse_command_spec_good() { |
|
130 | fn parse_command_spec_good() { | |
@@ -127,7 +167,66 b' mod tests {' | |||||
127 | assert!(parse_command_spec(Bytes::from_static(b"paper\0less")).is_err()); |
|
167 | assert!(parse_command_spec(Bytes::from_static(b"paper\0less")).is_err()); | |
128 | } |
|
168 | } | |
129 |
|
169 | |||
|
170 | #[test] | |||
|
171 | fn pack_env_vars_os_good() { | |||
|
172 | assert_eq!( | |||
|
173 | pack_env_vars_os(vec![] as Vec<(OsString, OsString)>), | |||
|
174 | Bytes::new() | |||
|
175 | ); | |||
|
176 | assert_eq!( | |||
|
177 | pack_env_vars_os(vec![os_string_pair_from(b"FOO", b"bar")]), | |||
|
178 | Bytes::from_static(b"FOO=bar") | |||
|
179 | ); | |||
|
180 | assert_eq!( | |||
|
181 | pack_env_vars_os(vec![ | |||
|
182 | os_string_pair_from(b"FOO", b""), | |||
|
183 | os_string_pair_from(b"BAR", b"baz") | |||
|
184 | ]), | |||
|
185 | Bytes::from_static(b"FOO=\0BAR=baz") | |||
|
186 | ); | |||
|
187 | } | |||
|
188 | ||||
|
189 | #[test] | |||
|
190 | fn pack_env_vars_os_large_key() { | |||
|
191 | let mut buf = vec![b'A'; INITIAL_PACKED_ENV_VARS_CAPACITY]; | |||
|
192 | let envs = vec![os_string_pair_from(&buf, b"")]; | |||
|
193 | buf.push(b'='); | |||
|
194 | assert_eq!(pack_env_vars_os(envs), Bytes::from(buf)); | |||
|
195 | } | |||
|
196 | ||||
|
197 | #[test] | |||
|
198 | fn pack_env_vars_os_large_value() { | |||
|
199 | let mut buf = vec![b'A', b'=']; | |||
|
200 | buf.resize(INITIAL_PACKED_ENV_VARS_CAPACITY + 1, b'a'); | |||
|
201 | let envs = vec![os_string_pair_from(&buf[..1], &buf[2..])]; | |||
|
202 | assert_eq!(pack_env_vars_os(envs), Bytes::from(buf)); | |||
|
203 | } | |||
|
204 | ||||
|
205 | #[test] | |||
|
206 | fn pack_env_vars_os_nul_eq() { | |||
|
207 | assert!(panic::catch_unwind(|| { | |||
|
208 | pack_env_vars_os(vec![os_string_pair_from(b"\0", b"")]) | |||
|
209 | }) | |||
|
210 | .is_err()); | |||
|
211 | assert!(panic::catch_unwind(|| { | |||
|
212 | pack_env_vars_os(vec![os_string_pair_from(b"FOO", b"\0bar")]) | |||
|
213 | }) | |||
|
214 | .is_err()); | |||
|
215 | assert!(panic::catch_unwind(|| { | |||
|
216 | pack_env_vars_os(vec![os_string_pair_from(b"FO=", b"bar")]) | |||
|
217 | }) | |||
|
218 | .is_err()); | |||
|
219 | assert_eq!( | |||
|
220 | pack_env_vars_os(vec![os_string_pair_from(b"FOO", b"=ba")]), | |||
|
221 | Bytes::from_static(b"FOO==ba") | |||
|
222 | ); | |||
|
223 | } | |||
|
224 | ||||
130 | fn os_string_from(s: &[u8]) -> OsString { |
|
225 | fn os_string_from(s: &[u8]) -> OsString { | |
131 | OsString::from_vec(s.to_vec()) |
|
226 | OsString::from_vec(s.to_vec()) | |
132 | } |
|
227 | } | |
|
228 | ||||
|
229 | fn os_string_pair_from(k: &[u8], v: &[u8]) -> (OsString, OsString) { | |||
|
230 | (os_string_from(k), os_string_from(v)) | |||
|
231 | } | |||
133 | } |
|
232 | } |
General Comments 0
You need to be logged in to leave comments.
Login now