ic-bn-lib 0.1.19

Internet Computer Boundary Nodes shared modules
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325

use std::{
    pin::{Pin, pin},
    sync::Mutex,
    sync::atomic::{AtomicBool, Ordering},
    task::{Context, Poll},
    time::Duration,
};

use axum::body::Body;
use bytes::{Buf, Bytes};
use futures_util::ready;
use http_body::{Body as HttpBody, Frame, SizeHint};
use http_body_util::{BodyExt, LengthLimitError, Limited};
use ic_bn_lib_common::types::http::Error;
use tokio::sync::{
    mpsc,
    oneshot::{self, Receiver, Sender},
};

use super::calc_headers_size;

/// Read the given body enforcing a size & time limit
pub async fn buffer_body<H: HttpBody + Send>(
    body: H,
    size_limit: usize,
    timeout: Duration,
) -> Result<Bytes, Error>
where
    <H as HttpBody>::Data: Buf + Send + Sync + 'static,
    <H as HttpBody>::Error: std::error::Error + Send + Sync + 'static,
{
    // Collect the request body up to the limit
    let body = tokio::time::timeout(timeout, Limited::new(body, size_limit).collect()).await;

    // Body reading timed out
    let Ok(body) = body else {
        return Err(Error::BodyTimedOut);
    };

    let body = body
        .map_err(|e| {
            // TODO improve the inferring somehow
            e.downcast_ref::<LengthLimitError>().map_or_else(
                || Error::BodyReadingFailed(e.to_string()),
                |_| Error::BodyTooBig,
            )
        })?
        .to_bytes();

    Ok(body)
}

/// Result of reading a body
pub type BodyResult = Result<u64, String>;

/// Wrapper that makes the provided body `Sync`
#[derive(Debug)]
pub struct SyncBody {
    inner: Mutex<Pin<Box<Body>>>,
}

impl SyncBody {
    /// Create a new `SyncBody`
    pub fn new(inner: Body) -> Self {
        Self {
            inner: Mutex::new(Box::pin(inner)),
        }
    }
}

impl HttpBody for SyncBody {
    type Data = Bytes;
    type Error = axum::Error;

    #[inline]
    fn poll_frame(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
        self.inner.lock().unwrap().as_mut().poll_frame(cx)
    }

    #[inline]
    fn is_end_stream(&self) -> bool {
        self.inner.lock().unwrap().as_ref().is_end_stream()
    }

    #[inline]
    fn size_hint(&self) -> SizeHint {
        self.inner.lock().unwrap().as_ref().size_hint()
    }
}

/// Body that notifies that it has finished by sending a value over the provided channel.
/// Use AtomicBool flag to make sure we notify only once.
pub struct NotifyingBody<D, E, S: Clone + Unpin> {
    inner: Pin<Box<dyn HttpBody<Data = D, Error = E> + Send + 'static>>,
    tx: mpsc::Sender<S>,
    sig: S,
    sent: AtomicBool,
}

impl<D, E, S: Clone + Unpin> NotifyingBody<D, E, S> {
    /// Create a new `NotifyingBody`
    pub fn new<B>(inner: B, tx: mpsc::Sender<S>, sig: S) -> Self
    where
        B: HttpBody<Data = D, Error = E> + Send + 'static,
        D: Buf,
    {
        Self {
            inner: Box::pin(inner),
            tx,
            sig,
            sent: AtomicBool::new(false),
        }
    }

    fn notify(&self) {
        if self
            .sent
            .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
            == Ok(false)
        {
            let _ = self.tx.try_send(self.sig.clone()).is_ok();
        }
    }
}

impl<D, E, S: Clone + Unpin> HttpBody for NotifyingBody<D, E, S>
where
    D: Buf,
    E: ToString,
{
    type Data = D;
    type Error = E;

    fn poll_frame(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
        let poll = ready!(pin!(&mut self.inner).poll_frame(cx));
        if poll.is_none() {
            self.notify();
        }

        Poll::Ready(poll)
    }

    fn size_hint(&self) -> SizeHint {
        self.inner.size_hint()
    }

    fn is_end_stream(&self) -> bool {
        let end = self.inner.is_end_stream();
        if end {
            self.notify();
        }

        end
    }
}

/// Body that counts the number of bytes streamed
pub struct CountingBody<D, E> {
    inner: Pin<Box<dyn HttpBody<Data = D, Error = E> + Send + 'static>>,
    tx: Option<Sender<BodyResult>>,
    expected_size: Option<u64>,
    bytes_sent: u64,
}

impl<D, E> CountingBody<D, E> {
    /// Create a new `CountingBody`
    pub fn new<B>(inner: B) -> (Self, Receiver<BodyResult>)
    where
        B: HttpBody<Data = D, Error = E> + Send + 'static,
        D: Buf,
    {
        let expected_size = inner.size_hint().exact();
        let (tx, rx) = oneshot::channel();

        let mut body = Self {
            inner: Box::pin(inner),
            tx: Some(tx),
            expected_size,
            bytes_sent: 0,
        };

        // If the size is known and zero - finish now,
        // otherwise it won't be called anywhere else
        if expected_size == Some(0) {
            body.finish(Ok(0));
        }

        (body, rx)
    }

    fn finish(&mut self, res: Result<u64, String>) {
        if let Some(v) = self.tx.take() {
            let _ = v.send(res);
        }
    }
}

impl<D, E> HttpBody for CountingBody<D, E>
where
    D: Buf,
    E: ToString,
{
    type Data = D;
    type Error = E;

    fn poll_frame(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
        let poll = ready!(pin!(&mut self.inner).poll_frame(cx));

        match &poll {
            // There is still some data available
            Some(v) => match v {
                Ok(buf) => {
                    // Normal data frame
                    if buf.is_data() {
                        self.bytes_sent += buf.data_ref().unwrap().remaining() as u64;
                    } else if buf.is_trailers() {
                        // Trailers are very uncommon, for the sake of completeness
                        self.bytes_sent += calc_headers_size(buf.trailers_ref().unwrap()) as u64;
                    }

                    // Check if we already got what was expected
                    if Some(self.bytes_sent) >= self.expected_size {
                        // Make borrow checker happy
                        let x = self.bytes_sent;
                        self.finish(Ok(x));
                    }
                }

                // Error occured
                Err(e) => {
                    self.finish(Err(e.to_string()));
                }
            },

            // Nothing left
            None => {
                // Make borrow checker happy
                let x = self.bytes_sent;
                self.finish(Ok(x));
            }
        }

        Poll::Ready(poll)
    }

    fn size_hint(&self) -> SizeHint {
        self.inner.size_hint()
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[tokio::test]
    async fn test_counting_body_stream() {
        let data = b"foobarblahblahfoobarblahblahfoobarblahblahfoobarblahblahfoobarbl\
        ahblahfoobarblahblahfoobarblahblahfoobarblahblahfoobarblahblahfoobarblahblahfoobarblahbla\
        hfoobarblahblahfoobarblahblahfoobarblahblahfoobarblahblahfoobarblahblahfoobarblahblahfoob\
        arblahblahfoobarblahblahfoobarblahblahfoobarblahblahfoobarblahblahfoobarblahblahfoobarbla\
        blahfoobarblahblah";

        let stream = tokio_util::io::ReaderStream::new(&data[..]);
        let body = Body::from_stream(stream);

        let (body, rx) = CountingBody::new(body);

        // Check that the body streams the same data back
        let body = body.collect().await.unwrap().to_bytes().to_vec();
        assert_eq!(body, data);

        // Check that the counting body got right number
        let size = rx.await.unwrap().unwrap();
        assert_eq!(size, data.len() as u64);
    }

    #[tokio::test]
    async fn test_counting_body_full() {
        let data = vec![0; 512];
        let buf = Bytes::from_iter(data.clone());
        let body = http_body_util::Full::new(buf);

        let (body, rx) = CountingBody::new(body);

        // Check that the body streams the same data back
        let body = body.collect().await.unwrap().to_bytes().to_vec();
        assert_eq!(body, data);

        // Check that the counting body got right number
        let size = rx.await.unwrap().unwrap();
        assert_eq!(size, data.len() as u64);
    }

    #[tokio::test]
    async fn test_notifying_body() {
        let data = b"foobarblahblahfoobarblahblahfoobarblahblahfoobarblahblahfoobarbl\
        ahblahfoobarblahblahfoobarblahblahfoobarblahblahfoobarblahblahfoobarblahblahfoobarblahbla\
        hfoobarblahblahfoobarblahblahfoobarblahblahfoobarblahblahfoobarblahblahfoobarblahblahfoob\
        arblahblahfoobarblahblahfoobarblahblahfoobarblahblahfoobarblahblahfoobarblahblahfoobarbla\
        blahfoobarblahblah";

        let stream = tokio_util::io::ReaderStream::new(&data[..]);
        let body = Body::from_stream(stream);

        let sig = 357;
        let (tx, mut rx) = mpsc::channel(10);
        let body = NotifyingBody::new(body, tx, sig);

        // Check that the body streams the same data back
        let body = body.collect().await.unwrap().to_bytes().to_vec();
        assert_eq!(body, data);

        // Make sure we're notified
        assert_eq!(sig, rx.recv().await.unwrap());
    }
}