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::{ensure, 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 impl<T> LengthPrefixed<T> {
36 /// Consumes the `LengthPrefixed` instance, returning the wrapped value.
into_inner(self) -> T37 pub fn into_inner(self) -> T {
38 self.inner
39 }
40 }
41
42 pub trait BytesExt {
read<T: ReadFromBytes>(&mut self) -> Result<T>43 fn read<T: ReadFromBytes>(&mut self) -> Result<T>;
44 }
45
46 impl BytesExt for Bytes {
read<T: ReadFromBytes>(&mut self) -> Result<T>47 fn read<T: ReadFromBytes>(&mut self) -> Result<T> {
48 T::read_from_bytes(self)
49 }
50 }
51
52 pub trait ReadFromBytes {
read_from_bytes(buf: &mut Bytes) -> Result<Self> where Self: Sized53 fn read_from_bytes(buf: &mut Bytes) -> Result<Self>
54 where
55 Self: Sized;
56 }
57
58 impl ReadFromBytes for u32 {
read_from_bytes(buf: &mut Bytes) -> Result<Self>59 fn read_from_bytes(buf: &mut Bytes) -> Result<Self> {
60 Ok(buf.get_u32_le())
61 }
62 }
63
64 impl<T: ReadFromBytes> ReadFromBytes for Vec<T> {
read_from_bytes(buf: &mut Bytes) -> Result<Self>65 fn read_from_bytes(buf: &mut Bytes) -> Result<Self> {
66 let mut result = vec![];
67 while buf.has_remaining() {
68 result.push(buf.read()?);
69 }
70 Ok(result)
71 }
72 }
73
74 impl<T: ReadFromBytes> ReadFromBytes for LengthPrefixed<T> {
read_from_bytes(buf: &mut Bytes) -> Result<Self>75 fn read_from_bytes(buf: &mut Bytes) -> Result<Self> {
76 let mut inner = read_length_prefixed_slice(buf)?;
77 let inner = inner.read()?;
78 Ok(LengthPrefixed { inner })
79 }
80 }
81
82 impl ReadFromBytes for Bytes {
read_from_bytes(buf: &mut Bytes) -> Result<Self>83 fn read_from_bytes(buf: &mut Bytes) -> Result<Self> {
84 Ok(buf.slice(..))
85 }
86 }
87
read_length_prefixed_slice(buf: &mut Bytes) -> Result<Bytes>88 fn read_length_prefixed_slice(buf: &mut Bytes) -> Result<Bytes> {
89 ensure!(
90 buf.remaining() >= 4,
91 "Remaining buffer too short to contain length of length-prefixed field. Remaining: {}",
92 buf.remaining()
93 );
94 let len = buf.get_u32_le() as usize;
95 ensure!(
96 buf.remaining() >= len,
97 "length-prefixed field longer than remaining buffer. Field length: {}, remaining: {}",
98 len,
99 buf.remaining()
100 );
101 Ok(buf.split_to(len))
102 }
103
104 #[cfg(test)]
105 mod tests {
106 use super::*;
107 use bytes::{BufMut, BytesMut};
108
109 #[test]
test_read_length_prefixed_slice()110 fn test_read_length_prefixed_slice() {
111 let data = b"hello world";
112 let mut b = BytesMut::new();
113 b.put_u32_le(data.len() as u32);
114 b.put_slice(data);
115 let mut slice = b.freeze();
116 let res = read_length_prefixed_slice(&mut slice);
117 assert!(res.is_ok());
118 assert_eq!(data, res.ok().unwrap().as_ref());
119 }
120 }
121