1 /*
2 * Copyright (C) 2021 The Android Open Source Project
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 //! Provides extension methods Bytes::read<T>(), which calls back ReadFromBytes::read_from_byte()
18
19 use anyhow::{bail, Result};
20 use bytes::{Buf, Bytes};
21 use std::ops::Deref;
22
23 #[derive(Clone, Debug)]
24 pub struct LengthPrefixed<T> {
25 inner: T,
26 }
27
28 impl<T> Deref for LengthPrefixed<T> {
29 type Target = T;
deref(&self) -> &Self::Target30 fn deref(&self) -> &Self::Target {
31 &self.inner
32 }
33 }
34
35 pub trait BytesExt {
read<T: ReadFromBytes>(&mut self) -> Result<T>36 fn read<T: ReadFromBytes>(&mut self) -> Result<T>;
37 }
38
39 impl BytesExt for Bytes {
read<T: ReadFromBytes>(&mut self) -> Result<T>40 fn read<T: ReadFromBytes>(&mut self) -> Result<T> {
41 T::read_from_bytes(self)
42 }
43 }
44
45 pub trait ReadFromBytes {
read_from_bytes(buf: &mut Bytes) -> Result<Self> where Self: Sized46 fn read_from_bytes(buf: &mut Bytes) -> Result<Self>
47 where
48 Self: Sized;
49 }
50
51 impl ReadFromBytes for u32 {
read_from_bytes(buf: &mut Bytes) -> Result<Self>52 fn read_from_bytes(buf: &mut Bytes) -> Result<Self> {
53 Ok(buf.get_u32_le())
54 }
55 }
56
57 impl<T: ReadFromBytes> ReadFromBytes for Vec<T> {
read_from_bytes(buf: &mut Bytes) -> Result<Self>58 fn read_from_bytes(buf: &mut Bytes) -> Result<Self> {
59 let mut result = vec![];
60 while buf.has_remaining() {
61 result.push(buf.read()?);
62 }
63 Ok(result)
64 }
65 }
66
67 impl<T: ReadFromBytes> ReadFromBytes for LengthPrefixed<T> {
read_from_bytes(buf: &mut Bytes) -> Result<Self>68 fn read_from_bytes(buf: &mut Bytes) -> Result<Self> {
69 let mut inner = read_length_prefixed_slice(buf)?;
70 let inner = inner.read()?;
71 Ok(LengthPrefixed { inner })
72 }
73 }
74
75 impl ReadFromBytes for Bytes {
read_from_bytes(buf: &mut Bytes) -> Result<Self>76 fn read_from_bytes(buf: &mut Bytes) -> Result<Self> {
77 Ok(buf.slice(..))
78 }
79 }
80
read_length_prefixed_slice(buf: &mut Bytes) -> Result<Bytes>81 fn read_length_prefixed_slice(buf: &mut Bytes) -> Result<Bytes> {
82 if buf.remaining() < 4 {
83 bail!(
84 "Remaining buffer too short to contain length of length-prefixed field. Remaining: {}",
85 buf.remaining()
86 );
87 }
88 let len = buf.get_u32_le() as usize;
89 if len > buf.remaining() {
90 bail!(
91 "length-prefixed field longer than remaining buffer. Field length: {}, remaining: {}",
92 len,
93 buf.remaining()
94 );
95 }
96 Ok(buf.split_to(len))
97 }
98
99 #[cfg(test)]
100 mod tests {
101 use super::*;
102 use bytes::{BufMut, BytesMut};
103
104 #[test]
test_read_length_prefixed_slice()105 fn test_read_length_prefixed_slice() {
106 let data = b"hello world";
107 let mut b = BytesMut::new();
108 b.put_u32_le(data.len() as u32);
109 b.put_slice(data);
110 let mut slice = b.freeze();
111 let res = read_length_prefixed_slice(&mut slice);
112 assert!(res.is_ok());
113 assert_eq!(data, res.ok().unwrap().as_ref());
114 }
115 }
116