1#[cfg(not(feature = "std"))]
4use alloc::vec::Vec;
5use bytes::{
6 Buf as _, Bytes, BytesMut, TryGetError,
7 buf::{Chain, Take},
8};
9use core::any::Any;
10
11pub trait Buf: bytes::Buf {
29 fn take_owned<T: OwnedBytes>(&mut self, len: usize) -> T {
38 T::from_bytes(self.copy_to_bytes(len))
39 }
40
41 fn try_take_owned<T: OwnedBytes>(&mut self, len: usize) -> Result<T, TryGetError> {
47 let available = self.remaining();
48 if available < len {
49 return Err(TryGetError {
50 requested: len,
51 available,
52 });
53 }
54 Ok(self.take_owned(len))
55 }
56}
57
58pub trait OwnedBytes: Sized + 'static {
64 fn from_bytes(bytes: Bytes) -> Self;
66
67 fn claim<V: Any>(view: V) -> Result<Self, V> {
71 let mut slot = Some(view);
73 let any: &mut dyn Any = &mut slot;
74 let claimed = any.downcast_mut::<Option<Self>>().and_then(Option::take);
75 claimed.ok_or_else(|| slot.expect("the slot must hold the view"))
76 }
77}
78
79impl OwnedBytes for Bytes {
80 #[inline]
81 fn from_bytes(bytes: Bytes) -> Self {
82 bytes
83 }
84}
85
86impl Buf for Bytes {}
87impl Buf for BytesMut {}
88impl<B: Buf + ?Sized> Buf for &mut B {
89 #[inline]
90 fn take_owned<T: OwnedBytes>(&mut self, len: usize) -> T {
91 (**self).take_owned(len)
92 }
93}
94
95impl<B: Buf> Buf for Take<B> {
96 #[inline]
98 fn take_owned<T: OwnedBytes>(&mut self, len: usize) -> T {
99 let remaining = self.remaining();
100 assert!(
101 len <= remaining,
102 "len must not exceed the remaining bytes: {len} > {remaining}"
103 );
104 let taken = self.get_mut().take_owned(len);
105 self.set_limit(self.limit() - len);
106 taken
107 }
108}
109
110impl<A: Buf, B: Buf> Buf for Chain<A, B> {
111 #[inline]
114 fn take_owned<T: OwnedBytes>(&mut self, len: usize) -> T {
115 let first = self.first_ref().remaining();
116 if first >= len {
117 self.first_mut().take_owned(len)
118 } else if first == 0 {
119 self.last_mut().take_owned(len)
120 } else {
121 T::from_bytes(self.copy_to_bytes(len))
122 }
123 }
124}
125
126pub trait Input: Into<Self::Buf> {
153 type Buf: Buf;
155}
156
157impl<B: Buf> Input for B {
158 type Buf = B;
159}
160
161impl Input for Vec<u8> {
162 type Buf = Bytes;
163}
164
165#[derive(Clone, Copy, Debug)]
201pub struct Copying<'a>(
202 pub &'a [u8],
204);
205
206impl bytes::Buf for Copying<'_> {
207 #[inline]
208 fn remaining(&self) -> usize {
209 self.0.remaining()
210 }
211
212 #[inline]
213 fn chunk(&self) -> &[u8] {
214 self.0.chunk()
215 }
216
217 #[inline]
218 fn advance(&mut self, cnt: usize) {
219 self.0.advance(cnt);
220 }
221
222 #[inline]
223 fn copy_to_slice(&mut self, dst: &mut [u8]) {
224 self.0.copy_to_slice(dst);
225 }
226}
227
228impl Buf for Copying<'_> {}
229
230#[derive(Clone, Copy, Debug)]
250pub struct Viewing<'a> {
251 bytes: &'a Bytes,
252 pos: usize,
253}
254
255impl<'a> Viewing<'a> {
256 pub const fn new(bytes: &'a Bytes) -> Self {
258 Self { bytes, pos: 0 }
259 }
260}
261
262impl bytes::Buf for Viewing<'_> {
263 #[inline]
264 fn remaining(&self) -> usize {
265 self.bytes.len() - self.pos
266 }
267
268 #[inline]
269 fn chunk(&self) -> &[u8] {
270 &self.bytes[self.pos..]
271 }
272
273 #[inline]
274 fn advance(&mut self, cnt: usize) {
275 assert!(
276 cnt <= self.remaining(),
277 "advance must not exceed the remaining bytes"
278 );
279 self.pos += cnt;
280 }
281
282 #[inline]
283 fn copy_to_bytes(&mut self, len: usize) -> Bytes {
284 assert!(
285 len <= self.remaining(),
286 "length must not exceed the remaining bytes"
287 );
288 let bytes = self.bytes.slice(self.pos..self.pos + len);
289 self.pos += len;
290 bytes
291 }
292}
293
294impl Buf for Viewing<'_> {}
295
296#[cfg(test)]
297mod tests {
298 use super::*;
299 use crate::{Decode, DecodeExt, Encode, Read, ReadExt};
300
301 #[test]
304 fn test_take_owned_claims_matching_views() {
305 struct Reader(Bytes);
307
308 #[derive(Debug, PartialEq)]
310 struct Owned {
311 bytes: Bytes,
312 claimed: bool,
313 }
314
315 impl bytes::Buf for Reader {
316 fn remaining(&self) -> usize {
317 self.0.remaining()
318 }
319 fn chunk(&self) -> &[u8] {
320 self.0.chunk()
321 }
322 fn advance(&mut self, cnt: usize) {
323 self.0.advance(cnt);
324 }
325 }
326
327 impl Buf for Reader {
328 fn take_owned<T: OwnedBytes>(&mut self, len: usize) -> T {
329 let view = Owned {
330 bytes: self.0.split_to(len),
331 claimed: true,
332 };
333 T::claim(view).unwrap_or_else(|view| T::from_bytes(view.bytes))
334 }
335 }
336
337 impl OwnedBytes for Owned {
338 fn from_bytes(bytes: Bytes) -> Self {
339 Self {
340 bytes,
341 claimed: false,
342 }
343 }
344 }
345
346 let mut reader = Reader(Bytes::from_static(b"hello world"));
347 let claimed: Owned = reader.take_owned(5);
348 assert_eq!(claimed.bytes, &b"hello"[..]);
349 assert!(claimed.claimed);
350 let bytes: Bytes = reader.take_owned(6);
351 assert_eq!(bytes, &b" world"[..]);
352
353 let mut plain = Bytes::from_static(b"hello");
355 let built: Owned = plain.take_owned(5);
356 assert_eq!(built.bytes, &b"hello"[..]);
357 assert!(!built.claimed);
358 }
359
360 #[test]
363 fn test_take_owned_through_adapters() {
364 struct Reader(Bytes);
365
366 #[derive(Debug, PartialEq)]
367 struct Owned(Bytes, bool);
368
369 impl bytes::Buf for Reader {
370 fn remaining(&self) -> usize {
371 self.0.remaining()
372 }
373 fn chunk(&self) -> &[u8] {
374 self.0.chunk()
375 }
376 fn advance(&mut self, cnt: usize) {
377 self.0.advance(cnt);
378 }
379 }
380
381 impl Buf for Reader {
382 fn take_owned<T: OwnedBytes>(&mut self, len: usize) -> T {
383 T::claim(Owned(self.0.split_to(len), true))
384 .unwrap_or_else(|view| T::from_bytes(view.0))
385 }
386 }
387
388 impl OwnedBytes for Owned {
389 fn from_bytes(bytes: Bytes) -> Self {
390 Self(bytes, false)
391 }
392 }
393
394 let mut reader = Reader(Bytes::from_static(b"hello world"));
395 let claimed: Owned = (&mut &mut reader).take_owned(5);
396 assert_eq!(claimed, Owned(Bytes::from_static(b"hello"), true));
397
398 let mut limited = bytes::Buf::take(&mut reader, 4);
399 let claimed: Owned = limited.take_owned(3);
400 assert_eq!(claimed, Owned(Bytes::from_static(b" wo"), true));
401 assert_eq!(limited.limit(), 1);
402 assert_eq!(limited.remaining(), 1);
403 let claimed: Owned = limited.try_take_owned(1).unwrap();
404 assert_eq!(claimed, Owned(Bytes::from_static(b"r"), true));
405 assert_eq!(limited.limit(), 0);
406
407 let mut chained =
410 Reader(Bytes::from_static(b"ab")).chain(Reader(Bytes::from_static(b"cd")));
411 let claimed: Owned = chained.take_owned(1);
412 assert_eq!(claimed, Owned(Bytes::from_static(b"a"), true));
413 let spanning: Owned = chained.take_owned(2);
414 assert_eq!(spanning, Owned(Bytes::from_static(b"bc"), false));
415 let claimed: Owned = chained.take_owned(1);
416 assert_eq!(claimed, Owned(Bytes::from_static(b"d"), true));
417 assert!(!chained.has_remaining());
418 }
419
420 #[test]
424 fn test_try_take_owned_reports_shortfall() {
425 let mut input = Bytes::from_static(b"hello");
426 let payload = input.as_ptr();
427 let err = input.try_take_owned::<Bytes>(6).unwrap_err();
428 assert_eq!((err.requested, err.available), (6, 5));
429 assert_eq!(input.remaining(), 5);
430 let taken: Bytes = input.try_take_owned(5).unwrap();
431 assert_eq!(taken.as_ptr(), payload);
432 assert!(!input.has_remaining());
433
434 let mut inner = Bytes::from_static(b"abc");
436 let mut limited = bytes::Buf::take(&mut inner, 10);
437 let err = limited.try_take_owned::<Bytes>(5).unwrap_err();
438 assert_eq!((err.requested, err.available), (5, 3));
439 assert_eq!(limited.remaining(), 3);
440 assert_eq!(limited.limit(), 10);
441
442 let mut inner = Bytes::from_static(b"abcdefgh");
444 let mut limited = bytes::Buf::take(&mut inner, 4);
445 let err = limited.try_take_owned::<Bytes>(5).unwrap_err();
446 assert_eq!((err.requested, err.available), (5, 4));
447 assert_eq!(limited.remaining(), 4);
448 assert_eq!(inner.remaining(), 8);
449 }
450
451 #[test]
454 #[should_panic(expected = "len must not exceed the remaining bytes: 5 > 3")]
455 fn test_take_owned_past_inner_end() {
456 let mut inner = Bytes::from_static(b"abc");
457 let mut limited = bytes::Buf::take(&mut inner, 10);
458 let _: Bytes = limited.take_owned(5);
459 }
460
461 #[test]
462 fn test_vec_preserves_payload() {
463 let encoded = Bytes::from_static(b"hello").encode().to_vec();
464 let payload = encoded[1..].as_ptr();
465 let decoded = Bytes::decode_cfg(encoded, &(..).into()).unwrap();
466 assert_eq!(decoded, b"hello"[..]);
467 assert_eq!(decoded.as_ptr(), payload);
468 }
469
470 #[test]
471 fn test_owned_buffers_preserve_payload() {
472 let encoded = Bytes::from_static(b"hello").encode();
473 let payload = encoded[1..].as_ptr();
474 let cfg = (..).into();
475 assert_eq!(
476 Bytes::decode_cfg(encoded.clone(), &cfg).unwrap().as_ptr(),
477 payload
478 );
479 assert_eq!(
480 Bytes::decode_cfg(encoded.clone().take(encoded.len()), &cfg)
481 .unwrap()
482 .as_ptr(),
483 payload
484 );
485 assert_eq!(
486 Bytes::decode_cfg(encoded.clone().chain(Bytes::new()), &cfg)
487 .unwrap()
488 .as_ptr(),
489 payload
490 );
491 let mut input = encoded;
492 assert_eq!(Bytes::read_cfg(&mut input, &cfg).unwrap().as_ptr(), payload);
493 assert_eq!(input.remaining(), 0);
494
495 let encoded = BytesMut::from(&b"\x05hello"[..]);
496 let payload = encoded[1..].as_ptr();
497 assert_eq!(Bytes::decode_cfg(encoded, &cfg).unwrap().as_ptr(), payload);
498 }
499
500 #[test]
501 fn test_copying_scratch() {
502 let scratch = [0, 0, 0, 7];
503 let mut input = Copying(&scratch);
504 assert_eq!(input.chunk().as_ptr(), scratch.as_ptr());
505 assert_eq!(u32::read(&mut input).unwrap(), 7);
506 assert_eq!(input.remaining(), 0);
507 assert_eq!(u32::decode(Copying(&scratch)).unwrap(), 7);
508 }
509
510 #[test]
511 fn test_copying_payload() {
512 let encoded = Bytes::from_static(b"hello").encode();
513 let decoded = Bytes::decode_cfg(Copying(&encoded), &(..).into()).unwrap();
514 assert_eq!(decoded, b"hello"[..]);
515 assert_ne!(decoded.as_ptr(), encoded[1..].as_ptr());
516 }
517
518 #[test]
520 fn test_viewing_scalar_reads() {
521 let encoded = Bytes::copy_from_slice(b"\x00\x00\x00\x07ab");
522 let mut input = Viewing::new(&encoded);
523 assert!(encoded.is_unique());
524 assert_eq!(input.chunk().as_ptr(), encoded.as_ptr());
525 assert_eq!(u32::read(&mut input).unwrap(), 7);
526
527 let mut other = input;
528 assert_eq!(<[u8; 2]>::read(&mut input).unwrap(), *b"ab");
529 assert_eq!(other.remaining(), 2);
530 assert_eq!(<[u8; 2]>::read(&mut other).unwrap(), *b"ab");
531 assert!(!input.has_remaining());
532 assert!(input.chunk().is_empty());
533 assert_eq!(encoded.as_ref(), b"\x00\x00\x00\x07ab");
534 assert!(encoded.is_unique());
535 }
536
537 #[test]
539 fn test_viewing_retained_fields() {
540 let (first, second) = {
541 let encoded = Bytes::copy_from_slice(b"\x05hello\x05world");
542 let mut input = Viewing::new(&encoded);
543 let mut limited = (&mut input).take(6);
544 let first = Bytes::read_cfg(&mut limited, &(..).into()).unwrap();
545 assert_eq!(first.as_ptr(), encoded[1..].as_ptr());
546 assert!(!limited.has_remaining());
547
548 let second = Bytes::decode_cfg(input, &(..).into()).unwrap();
549 assert_eq!(second.as_ptr(), encoded[7..].as_ptr());
550 (first, second)
551 };
552 assert_eq!(first, b"hello"[..]);
553 assert_eq!(second, b"world"[..]);
554 }
555
556 #[test]
558 fn test_viewing_empty_ranges() {
559 for encoded in [Bytes::new(), Bytes::from_static(b"hello")] {
560 let mut input = Viewing::new(&encoded);
561 input.advance(0);
562 assert!(input.copy_to_bytes(0).is_empty());
563 assert_eq!(input.remaining(), encoded.len());
564 input.advance(encoded.len());
565 input.advance(0);
566 assert!(input.copy_to_bytes(0).is_empty());
567 assert!(!input.has_remaining());
568 assert!(input.chunk().is_empty());
569 }
570 }
571
572 #[test]
574 fn test_viewing_invalid_input() {
575 for bytes in [&b""[..], b"\x05hi", b"\x01ab", b"\x80"] {
576 let encoded = Bytes::copy_from_slice(bytes);
577 assert!(Bytes::decode_cfg(Viewing::new(&encoded), &(..).into()).is_err());
578 }
579 let encoded = Bytes::from_static(b"\x00");
580 assert!(
581 Bytes::decode_cfg(Viewing::new(&encoded), &(..).into())
582 .unwrap()
583 .is_empty()
584 );
585 }
586
587 #[test]
589 #[should_panic(expected = "advance must not exceed the remaining bytes")]
590 fn test_viewing_advance_past_end() {
591 let encoded = Bytes::from_static(b"a");
592 let mut input = Viewing::new(&encoded);
593 input.advance(1);
594 input.advance(usize::MAX);
595 }
596
597 #[test]
599 #[should_panic(expected = "length must not exceed the remaining bytes")]
600 fn test_viewing_copy_past_end() {
601 let encoded = Bytes::from_static(b"a");
602 let mut input = Viewing::new(&encoded);
603 input.advance(1);
604 input.copy_to_bytes(usize::MAX);
605 }
606}