1 use crate::Error;
2 use std::collections::HashMap;
3 use std::convert::TryInto;
4 use std::ffi::{CString, OsString};
5 use wiggle::GuestPtr;
6 
7 #[derive(Debug, Eq, Hash, PartialEq)]
8 pub enum PendingString {
9     Bytes(Vec<u8>),
10     OsString(OsString),
11 }
12 
13 impl From<Vec<u8>> for PendingString {
14     fn from(bytes: Vec<u8>) -> Self {
15         Self::Bytes(bytes)
16     }
17 }
18 
19 impl From<OsString> for PendingString {
20     fn from(s: OsString) -> Self {
21         Self::OsString(s)
22     }
23 }
24 
25 impl PendingString {
26     pub fn into_string(self) -> Result<String, StringArrayError> {
27         let res = match self {
28             Self::Bytes(v) => String::from_utf8(v)?,
29             #[cfg(unix)]
30             Self::OsString(s) => {
31                 use std::os::unix::ffi::OsStringExt;
32                 String::from_utf8(s.into_vec())?
33             }
34             #[cfg(windows)]
35             Self::OsString(s) => {
36                 use std::os::windows::ffi::OsStrExt;
37                 let bytes: Vec<u16> = s.encode_wide().collect();
38                 String::from_utf16(&bytes)?
39             }
40         };
41         Ok(res)
42     }
43 }
44 
45 #[derive(Debug, thiserror::Error)]
46 pub enum StringArrayError {
47     /// Provided sequence of bytes contained an unexpected NUL byte.
48     #[error("provided sequence contained an unexpected NUL byte")]
49     Nul(#[from] std::ffi::NulError),
50     /// Too many elements: must fit into u32
51     #[error("too many elements")]
52     NumElements,
53     /// Element size: must fit into u32
54     #[error("element too big")]
55     ElemSize,
56     /// Cumulative element size: must fit into u32
57     #[error("cumulative element size too big")]
58     CumElemSize,
59     /// Provided sequence of bytes was not a valid UTF-8.
60     #[error("provided sequence is not valid UTF-8: {0}")]
61     InvalidUtf8(#[from] std::string::FromUtf8Error),
62     /// Provided sequence of bytes was not a valid UTF-16.
63     ///
64     /// This error is expected to only occur on Windows hosts.
65     #[error("provided sequence is not valid UTF-16: {0}")]
66     InvalidUtf16(#[from] std::string::FromUtf16Error),
67 }
68 
69 pub struct StringArray {
70     elems: Vec<CString>,
71     pub number_elements: u32,
72     pub cumulative_size: u32,
73 }
74 impl StringArray {
75     pub fn from_pending_vec(elems: Vec<PendingString>) -> Result<Self, StringArrayError> {
76         let elems = elems
77             .into_iter()
78             .map(|arg| arg.into_string())
79             .collect::<Result<Vec<String>, StringArrayError>>()?;
80         Self::from_strings(elems)
81     }
82     pub fn from_pending_map(
83         elems: HashMap<PendingString, PendingString>,
84     ) -> Result<Self, StringArrayError> {
85         let mut pairs = Vec::new();
86         for (k, v) in elems.into_iter() {
87             let mut pair = k.into_string()?;
88             pair.push('=');
89             pair.push_str(&v.into_string()?);
90             pairs.push(pair);
91         }
92         Self::from_strings(pairs)
93     }
94     pub fn from_strings(elems: Vec<String>) -> Result<Self, StringArrayError> {
95         let elems = elems
96             .into_iter()
97             .map(|s| CString::new(s))
98             .collect::<Result<Vec<CString>, _>>()?;
99         let number_elements = elems
100             .len()
101             .try_into()
102             .map_err(|_| StringArrayError::NumElements)?;
103         let mut cumulative_size: u32 = 0;
104         for elem in elems.iter() {
105             let elem_bytes = elem
106                 .as_bytes_with_nul()
107                 .len()
108                 .try_into()
109                 .map_err(|_| StringArrayError::ElemSize)?;
110             cumulative_size = cumulative_size
111                 .checked_add(elem_bytes)
112                 .ok_or(StringArrayError::CumElemSize)?;
113         }
114         Ok(Self {
115             elems,
116             number_elements,
117             cumulative_size,
118         })
119     }
120 
121     pub fn write_to_guest<'a>(
122         &self,
123         buffer: &GuestPtr<'a, u8>,
124         element_heads: &GuestPtr<'a, GuestPtr<'a, u8>>,
125     ) -> Result<(), Error> {
126         let element_heads = element_heads.as_array(self.number_elements);
127         let buffer = buffer.as_array(self.cumulative_size);
128         let mut cursor = 0;
129         for (elem, head) in self.elems.iter().zip(element_heads.iter()) {
130             let bytes = elem.as_bytes_with_nul();
131             let len: u32 = bytes.len().try_into()?;
132             let elem_buffer = buffer
133                 .get_range(cursor..(cursor + len))
134                 .ok_or(Error::Inval)?; // Elements don't fit in buffer provided
135             elem_buffer.copy_from_slice(bytes)?;
136             head?.write(
137                 elem_buffer
138                     .get(0)
139                     .expect("all elem buffers at least length 1"),
140             )?;
141             cursor += len;
142         }
143         Ok(())
144     }
145 }
146