capybara_core/protocol/http/frame/
query.rs

1use std::borrow::Cow;
2
3use memchr::memmem;
4use smallvec::SmallVec;
5
6#[derive(Debug)]
7pub struct Queries<'a>(&'a [u8]);
8
9impl Queries<'_> {
10    pub fn iter(&self) -> Iter<'_> {
11        Iter {
12            queries: self,
13            cur: 0,
14        }
15    }
16
17    pub fn as_bytes(&self) -> &[u8] {
18        self.0
19    }
20
21    pub fn as_str(&self) -> &str {
22        unsafe { std::str::from_utf8_unchecked(self.0) }
23    }
24
25    pub fn len(&self) -> usize {
26        self.iter().count()
27    }
28
29    pub fn is_empty(&self) -> bool {
30        self.iter().next().is_none()
31    }
32
33    pub fn get_last<K>(&self, key: K) -> Option<&[u8]>
34    where
35        K: AsRef<[u8]>,
36    {
37        let key = key.as_ref();
38
39        let mut result = None;
40
41        if key.is_empty() {
42            return result;
43        }
44
45        let finder = memmem::Finder::new(key);
46        for i in finder.find_iter(self.0) {
47            if i != 0 && self.0[i - 1] != b'&' {
48                continue;
49            }
50
51            let mut offset = i + key.len();
52
53            // EOF
54            if offset == self.0.len() {
55                result.replace(&b""[..]);
56                break;
57            }
58
59            // check next char:
60            //  - ...&foo&... => None
61            //  - ...&foo=&... => None
62            //  - ...&foo=xxx&... => Some(xxx)
63            match self.0[offset] {
64                b'&' => {
65                    result.replace(&b""[..]);
66                }
67                b'=' => {
68                    offset += 1;
69                    let mut end = None;
70                    for i in offset..self.0.len() {
71                        if self.0[i] == b'&' {
72                            end.replace(i);
73                            break;
74                        }
75                    }
76                    match end {
77                        None => {
78                            result.replace(&self.0[offset..]);
79                        }
80                        Some(j) => {
81                            result.replace(&self.0[offset..j]);
82                        }
83                    }
84                }
85                _ => (),
86            }
87        }
88
89        result
90    }
91
92    pub fn get_first<K>(&self, key: K) -> Option<&[u8]>
93    where
94        K: AsRef<[u8]>,
95    {
96        let key = key.as_ref();
97
98        if key.is_empty() {
99            return None;
100        }
101
102        let finder = memmem::Finder::new(key);
103        for i in finder.find_iter(self.0) {
104            if i != 0 && self.0[i - 1] != b'&' {
105                continue;
106            }
107
108            let mut offset = i + key.len();
109
110            // EOF
111            if offset == self.0.len() {
112                return Some(&b""[..]);
113            }
114
115            // check next char:
116            //  - ...&foo&... => None
117            //  - ...&foo=&... => None
118            //  - ...&foo=xxx&... => Some(xxx)
119            match self.0[offset] {
120                b'&' => {
121                    return Some(&b""[..]);
122                }
123                b'=' => {
124                    offset += 1;
125                    let mut end = None;
126                    for i in offset..self.0.len() {
127                        if self.0[i] == b'&' {
128                            end.replace(i);
129                            break;
130                        }
131                    }
132                    return match end {
133                        None => Some(&self.0[offset..]),
134                        Some(j) => Some(&self.0[offset..j]),
135                    };
136                }
137                _ => (),
138            }
139        }
140
141        None
142    }
143
144    pub fn get_all<K>(&self, key: K) -> SmallVec<[&[u8]; 4]>
145    where
146        K: AsRef<[u8]>,
147    {
148        let key = key.as_ref();
149
150        let mut result = SmallVec::<[&[u8]; 4]>::new();
151
152        if key.is_empty() {
153            return result;
154        }
155
156        let finder = memmem::Finder::new(key);
157        for i in finder.find_iter(self.0) {
158            if i != 0 && self.0[i - 1] != b'&' {
159                continue;
160            }
161
162            let mut offset = i + key.len();
163
164            // EOF
165            if offset == self.0.len() {
166                result.push(&b""[..]);
167                break;
168            }
169
170            // check next char:
171            //  - ...&foo&... => None
172            //  - ...&foo=&... => None
173            //  - ...&foo=xxx&... => Some(xxx)
174            match self.0[offset] {
175                b'&' => {
176                    result.push(&b""[..]);
177                }
178                b'=' => {
179                    offset += 1;
180                    let mut end = None;
181                    for i in offset..self.0.len() {
182                        if self.0[i] == b'&' {
183                            end.replace(i);
184                            break;
185                        }
186                    }
187                    match end {
188                        None => result.push(&self.0[offset..]),
189                        Some(j) => result.push(&self.0[offset..j]),
190                    }
191                }
192                _ => (),
193            }
194        }
195
196        result
197    }
198
199    pub fn get(&self, key: &str) -> Option<Query> {
200        if key.is_empty() {
201            return None;
202        }
203
204        let first = key.as_bytes().first().unwrap();
205        for next in self.iter() {
206            if next.0.is_empty() {
207                continue;
208            }
209
210            if !next.0.first().unwrap().eq_ignore_ascii_case(first) {
211                continue;
212            }
213
214            if next.key().eq_ignore_ascii_case(key) {
215                return Some(next);
216            }
217        }
218
219        None
220    }
221}
222
223impl<'a> From<&'a [u8]> for Queries<'a> {
224    fn from(value: &'a [u8]) -> Self {
225        Self(value)
226    }
227}
228
229#[derive(Debug)]
230pub struct Query<'a>(pub(crate) &'a [u8]);
231
232impl Query<'_> {
233    pub fn key(&self) -> Cow<str> {
234        match self.pos() {
235            Some(i) => String::from_utf8_lossy(&self.0[..i]),
236            None => String::from_utf8_lossy(self.0),
237        }
238    }
239
240    pub fn value_bytes(&self) -> Option<&[u8]> {
241        self.pos().map(|i| &self.0[i + 1..])
242    }
243
244    pub fn value(&self) -> Option<Cow<str>> {
245        self.pos()
246            .map(|i| String::from_utf8_lossy(&self.0[i + 1..]))
247    }
248
249    #[inline]
250    fn pos(&self) -> Option<usize> {
251        self.0.iter().position(|b| *b == b'=')
252    }
253
254    pub fn as_str(&self) -> &str {
255        unsafe { std::str::from_utf8_unchecked(self.0) }
256    }
257
258    pub fn as_bytes(&self) -> &[u8] {
259        self.0
260    }
261}
262
263impl<'a> AsRef<str> for Query<'a> {
264    fn as_ref(&self) -> &str {
265        unsafe { std::str::from_utf8_unchecked(self.0) }
266    }
267}
268
269impl<'a> AsRef<[u8]> for Query<'a> {
270    fn as_ref(&self) -> &[u8] {
271        self.0
272    }
273}
274
275impl<'a> From<&'a [u8]> for Query<'a> {
276    fn from(b: &'a [u8]) -> Self {
277        Self(b)
278    }
279}
280
281pub struct Iter<'a> {
282    queries: &'a Queries<'a>,
283    cur: isize,
284}
285
286impl<'a> Iterator for Iter<'a> {
287    type Item = Query<'a>;
288
289    fn next(&mut self) -> Option<Self::Item> {
290        // eof
291        if self.cur == -1 {
292            return None;
293        }
294
295        // overflow
296        if self.cur >= self.queries.0.len() as isize {
297            return None;
298        }
299
300        let data: &[u8] = &self.queries.0[self.cur as usize..];
301        match data.iter().position(|b| *b == b'&') {
302            Some(to) => {
303                self.cur += to as isize + 1;
304                let data = &data[..to];
305                if data.is_empty() {
306                    self.next()
307                } else {
308                    Some(Query(data))
309                }
310            }
311            None => {
312                self.cur = -1;
313                if data.is_empty() {
314                    self.next()
315                } else {
316                    Some(Query(data))
317                }
318            }
319        }
320    }
321}
322
323#[cfg(test)]
324mod queries_tests {
325    use super::*;
326
327    fn init() {
328        pretty_env_logger::try_init_timed().ok();
329    }
330
331    #[test]
332    fn queries() {
333        init();
334
335        let raw = b"foo=1&&x=2&details";
336        let queries = Queries::from(&raw[..]);
337
338        let x = queries.get("x");
339        assert!(x.is_some());
340        let x = x.unwrap();
341        assert_eq!("x", &x.key());
342        assert_eq!("2", &x.value().unwrap());
343
344        for q in queries.iter() {
345            let k = q.key();
346            let v = q.value();
347
348            info!("next query: |{}|{}|{:?}|", q.as_str(), k, v);
349
350            match k.as_ref() {
351                "foo" => assert_eq!(Some("1".into()), v),
352                "x" => assert_eq!(Some("2".into()), v),
353                "details" => assert!(v.is_none()),
354                k => {
355                    error!("bad key {}", k)
356                }
357            }
358        }
359    }
360
361    #[test]
362    fn test_weird_values() {
363        init();
364
365        let raw = b"x=a=1;b=2";
366        let queries = Queries::from(&raw[..]);
367    }
368
369    #[test]
370    fn test_values() {
371        init();
372
373        let raw = b"x=a&y=2&x=&x=b&z=3&x=c&x";
374        let queries = Queries::from(&raw[..]);
375
376        let found = queries.get_all("x");
377        assert_eq!(5, found.len());
378
379        for next in found.iter() {
380            info!("next: |{}|", String::from_utf8_lossy(next));
381        }
382
383        let mut it = found.into_iter();
384        assert_eq!(Some(&b"a"[..]), it.next());
385        assert_eq!(Some(&b""[..]), it.next());
386        assert_eq!(Some(&b"b"[..]), it.next());
387        assert_eq!(Some(&b"c"[..]), it.next());
388        assert_eq!(Some(&b""[..]), it.next());
389        assert!(it.next().is_none());
390
391        let first_x = queries.get_first("x");
392        assert_eq!(Some(&b"a"[..]), first_x);
393
394        let last_x = queries.get_last("x");
395        assert_eq!(Some(&b""[..]), last_x);
396
397        let y = queries.get_all("y");
398        assert_eq!(1, y.len());
399        let mut it = y.into_iter();
400        assert_eq!(Some(&b"2"[..]), it.next());
401        assert!(it.next().is_none());
402
403        let z = queries.get_all("z");
404        assert_eq!(1, z.len());
405        let mut it = z.into_iter();
406        assert_eq!(Some(&b"3"[..]), it.next());
407        assert!(it.next().is_none());
408
409        let last_z = queries.get_last("z");
410        assert_eq!(Some(&b"3"[..]), last_z);
411    }
412}
capybara_core/protocol/http/frame/query.rs

capybara_core/protocol/http/frame/
query.rs
