1 | use crate::adapter::StripBytes; |
2 | use crate::IsTerminal; |
3 | use crate::Lockable; |
4 | use crate::RawStream; |
5 | |
6 | /// Only pass printable data to the inner `Write` |
7 | #[derive (Debug)] |
8 | pub struct StripStream<S> { |
9 | raw: S, |
10 | state: StripBytes, |
11 | } |
12 | |
13 | impl<S> StripStream<S> |
14 | where |
15 | S: RawStream, |
16 | { |
17 | /// Only pass printable data to the inner `Write` |
18 | #[inline ] |
19 | pub fn new(raw: S) -> Self { |
20 | Self { |
21 | raw, |
22 | state: Default::default(), |
23 | } |
24 | } |
25 | |
26 | /// Get the wrapped [`RawStream`] |
27 | #[inline ] |
28 | pub fn into_inner(self) -> S { |
29 | self.raw |
30 | } |
31 | |
32 | #[inline ] |
33 | pub fn is_terminal(&self) -> bool { |
34 | self.raw.is_terminal() |
35 | } |
36 | } |
37 | |
38 | impl<S> IsTerminal for StripStream<S> |
39 | where |
40 | S: RawStream, |
41 | { |
42 | #[inline ] |
43 | fn is_terminal(&self) -> bool { |
44 | self.is_terminal() |
45 | } |
46 | } |
47 | |
48 | impl<S> std::io::Write for StripStream<S> |
49 | where |
50 | S: RawStream, |
51 | { |
52 | #[inline ] |
53 | fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> { |
54 | write(&mut self.raw, &mut self.state, buf) |
55 | } |
56 | |
57 | #[inline ] |
58 | fn flush(&mut self) -> std::io::Result<()> { |
59 | self.raw.flush() |
60 | } |
61 | |
62 | // Provide explicit implementations of trait methods |
63 | // - To reduce bookkeeping |
64 | // - Avoid acquiring / releasing locks in a loop |
65 | |
66 | #[inline ] |
67 | fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> { |
68 | write_all(&mut self.raw, &mut self.state, buf) |
69 | } |
70 | |
71 | // Not bothering with `write_fmt` as it just calls `write_all` |
72 | } |
73 | |
74 | fn write( |
75 | raw: &mut dyn std::io::Write, |
76 | state: &mut StripBytes, |
77 | buf: &[u8], |
78 | ) -> std::io::Result<usize> { |
79 | let initial_state: StripBytes = state.clone(); |
80 | |
81 | for printable: &[u8] in state.strip_next(bytes:buf) { |
82 | let possible: usize = printable.len(); |
83 | let written: usize = raw.write(buf:printable)?; |
84 | if possible != written { |
85 | let divergence: &[u8] = &printable[written..]; |
86 | let offset: usize = offset_to(total:buf, subslice:divergence); |
87 | let consumed: &[u8] = &buf[offset..]; |
88 | *state = initial_state; |
89 | state.strip_next(bytes:consumed).last(); |
90 | return Ok(offset); |
91 | } |
92 | } |
93 | Ok(buf.len()) |
94 | } |
95 | |
96 | fn write_all( |
97 | raw: &mut dyn std::io::Write, |
98 | state: &mut StripBytes, |
99 | buf: &[u8], |
100 | ) -> std::io::Result<()> { |
101 | for printable: &[u8] in state.strip_next(bytes:buf) { |
102 | raw.write_all(buf:printable)?; |
103 | } |
104 | Ok(()) |
105 | } |
106 | |
107 | #[inline ] |
108 | fn offset_to(total: &[u8], subslice: &[u8]) -> usize { |
109 | let total: *const u8 = total.as_ptr(); |
110 | let subslice: *const u8 = subslice.as_ptr(); |
111 | |
112 | debug_assert!( |
113 | total <= subslice, |
114 | "`Offset::offset_to` only accepts slices of `self`" |
115 | ); |
116 | subslice as usize - total as usize |
117 | } |
118 | |
119 | impl<S> Lockable for StripStream<S> |
120 | where |
121 | S: Lockable, |
122 | { |
123 | type Locked = StripStream<<S as Lockable>::Locked>; |
124 | |
125 | #[inline ] |
126 | fn lock(self) -> Self::Locked { |
127 | Self::Locked { |
128 | raw: self.raw.lock(), |
129 | state: self.state, |
130 | } |
131 | } |
132 | } |
133 | |
134 | #[cfg (test)] |
135 | mod test { |
136 | use super::*; |
137 | use proptest::prelude::*; |
138 | use std::io::Write as _; |
139 | |
140 | proptest! { |
141 | #[test] |
142 | #[cfg_attr(miri, ignore)] // See https://github.com/AltSysrq/proptest/issues/253 |
143 | fn write_all_no_escapes(s in " \\PC*" ) { |
144 | let buffer = crate::Buffer::new(); |
145 | let mut stream = StripStream::new(buffer); |
146 | stream.write_all(s.as_bytes()).unwrap(); |
147 | let buffer = stream.into_inner(); |
148 | let actual = std::str::from_utf8(buffer.as_ref()).unwrap(); |
149 | assert_eq!(s, actual); |
150 | } |
151 | |
152 | #[test] |
153 | #[cfg_attr(miri, ignore)] // See https://github.com/AltSysrq/proptest/issues/253 |
154 | fn write_byte_no_escapes(s in " \\PC*" ) { |
155 | let buffer = crate::Buffer::new(); |
156 | let mut stream = StripStream::new(buffer); |
157 | for byte in s.as_bytes() { |
158 | stream.write_all(&[*byte]).unwrap(); |
159 | } |
160 | let buffer = stream.into_inner(); |
161 | let actual = std::str::from_utf8(buffer.as_ref()).unwrap(); |
162 | assert_eq!(s, actual); |
163 | } |
164 | |
165 | #[test] |
166 | #[cfg_attr(miri, ignore)] // See https://github.com/AltSysrq/proptest/issues/253 |
167 | fn write_all_random(s in any::<Vec<u8>>()) { |
168 | let buffer = crate::Buffer::new(); |
169 | let mut stream = StripStream::new(buffer); |
170 | stream.write_all(s.as_slice()).unwrap(); |
171 | let buffer = stream.into_inner(); |
172 | if let Ok(actual) = std::str::from_utf8(buffer.as_ref()) { |
173 | for char in actual.chars() { |
174 | assert!(!char.is_ascii() || !char.is_control() || char.is_ascii_whitespace(), "{:?} -> {:?}: {:?}" , String::from_utf8_lossy(&s), actual, char); |
175 | } |
176 | } |
177 | } |
178 | |
179 | #[test] |
180 | #[cfg_attr(miri, ignore)] // See https://github.com/AltSysrq/proptest/issues/253 |
181 | fn write_byte_random(s in any::<Vec<u8>>()) { |
182 | let buffer = crate::Buffer::new(); |
183 | let mut stream = StripStream::new(buffer); |
184 | for byte in s.as_slice() { |
185 | stream.write_all(&[*byte]).unwrap(); |
186 | } |
187 | let buffer = stream.into_inner(); |
188 | if let Ok(actual) = std::str::from_utf8(buffer.as_ref()) { |
189 | for char in actual.chars() { |
190 | assert!(!char.is_ascii() || !char.is_control() || char.is_ascii_whitespace(), "{:?} -> {:?}: {:?}" , String::from_utf8_lossy(&s), actual, char); |
191 | } |
192 | } |
193 | } |
194 | } |
195 | } |
196 | |