1 | use crate::{Body, SizeHint}; |
2 | use bytes::{Buf, Bytes}; |
3 | use http::HeaderMap; |
4 | use pin_project_lite::pin_project; |
5 | use std::borrow::Cow; |
6 | use std::convert::{Infallible, TryFrom}; |
7 | use std::pin::Pin; |
8 | use std::task::{Context, Poll}; |
9 | |
10 | pin_project! { |
11 | /// A body that consists of a single chunk. |
12 | #[derive(Clone, Copy, Debug)] |
13 | pub struct Full<D> { |
14 | data: Option<D>, |
15 | } |
16 | } |
17 | |
18 | impl<D> Full<D> |
19 | where |
20 | D: Buf, |
21 | { |
22 | /// Create a new `Full`. |
23 | pub fn new(data: D) -> Self { |
24 | let data: Option = if data.has_remaining() { |
25 | Some(data) |
26 | } else { |
27 | None |
28 | }; |
29 | Full { data } |
30 | } |
31 | } |
32 | |
33 | impl<D> Body for Full<D> |
34 | where |
35 | D: Buf, |
36 | { |
37 | type Data = D; |
38 | type Error = Infallible; |
39 | |
40 | fn poll_data( |
41 | mut self: Pin<&mut Self>, |
42 | _cx: &mut Context<'_>, |
43 | ) -> Poll<Option<Result<D, Self::Error>>> { |
44 | Poll::Ready(self.data.take().map(Ok)) |
45 | } |
46 | |
47 | fn poll_trailers( |
48 | self: Pin<&mut Self>, |
49 | _cx: &mut Context<'_>, |
50 | ) -> Poll<Result<Option<HeaderMap>, Self::Error>> { |
51 | Poll::Ready(Ok(None)) |
52 | } |
53 | |
54 | fn is_end_stream(&self) -> bool { |
55 | self.data.is_none() |
56 | } |
57 | |
58 | fn size_hint(&self) -> SizeHint { |
59 | self.data |
60 | .as_ref() |
61 | .map(|data| SizeHint::with_exact(u64::try_from(data.remaining()).unwrap())) |
62 | .unwrap_or_else(|| SizeHint::with_exact(0)) |
63 | } |
64 | } |
65 | |
66 | impl<D> Default for Full<D> |
67 | where |
68 | D: Buf, |
69 | { |
70 | /// Create an empty `Full`. |
71 | fn default() -> Self { |
72 | Full { data: None } |
73 | } |
74 | } |
75 | |
76 | impl<D> From<Bytes> for Full<D> |
77 | where |
78 | D: Buf + From<Bytes>, |
79 | { |
80 | fn from(bytes: Bytes) -> Self { |
81 | Full::new(D::from(bytes)) |
82 | } |
83 | } |
84 | |
85 | impl<D> From<Vec<u8>> for Full<D> |
86 | where |
87 | D: Buf + From<Vec<u8>>, |
88 | { |
89 | fn from(vec: Vec<u8>) -> Self { |
90 | Full::new(D::from(vec)) |
91 | } |
92 | } |
93 | |
94 | impl<D> From<&'static [u8]> for Full<D> |
95 | where |
96 | D: Buf + From<&'static [u8]>, |
97 | { |
98 | fn from(slice: &'static [u8]) -> Self { |
99 | Full::new(D::from(slice)) |
100 | } |
101 | } |
102 | |
103 | impl<D, B> From<Cow<'static, B>> for Full<D> |
104 | where |
105 | D: Buf + From<&'static B> + From<B::Owned>, |
106 | B: ToOwned + ?Sized, |
107 | { |
108 | fn from(cow: Cow<'static, B>) -> Self { |
109 | match cow { |
110 | Cow::Borrowed(b: &B) => Full::new(D::from(b)), |
111 | Cow::Owned(o: ::Owned) => Full::new(D::from(o)), |
112 | } |
113 | } |
114 | } |
115 | |
116 | impl<D> From<String> for Full<D> |
117 | where |
118 | D: Buf + From<String>, |
119 | { |
120 | fn from(s: String) -> Self { |
121 | Full::new(D::from(s)) |
122 | } |
123 | } |
124 | |
125 | impl<D> From<&'static str> for Full<D> |
126 | where |
127 | D: Buf + From<&'static str>, |
128 | { |
129 | fn from(slice: &'static str) -> Self { |
130 | Full::new(D::from(slice)) |
131 | } |
132 | } |
133 | |
134 | #[cfg (test)] |
135 | mod tests { |
136 | use super::*; |
137 | |
138 | #[tokio::test] |
139 | async fn full_returns_some() { |
140 | let mut full = Full::new(&b"hello" [..]); |
141 | assert_eq!(full.size_hint().exact(), Some(b"hello" .len() as u64)); |
142 | assert_eq!(full.data().await, Some(Ok(&b"hello" [..]))); |
143 | assert!(full.data().await.is_none()); |
144 | } |
145 | |
146 | #[tokio::test] |
147 | async fn empty_full_returns_none() { |
148 | assert!(Full::<&[u8]>::default().data().await.is_none()); |
149 | assert!(Full::new(&b"" [..]).data().await.is_none()); |
150 | } |
151 | } |
152 | |