proxide 0.2.3

Proxide is a debugging proxy for capturing HTTP/2 and gRPC traffic. 
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

use std::io::Result;
use std::pin::Pin;
use std::task::Context;
use std::task::Poll;
use tokio::io::{split, AsyncRead, AsyncWrite, ReadBuf, ReadHalf, WriteHalf};

pub struct PrefixedRead<S>
{
    prefix: Option<Vec<u8>>,
    read: ReadHalf<S>,
}

pub struct PrefixedStream<S>
{
    read: PrefixedRead<S>,
    write: WriteHalf<S>,
}

impl<S> PrefixedStream<S>
where
    S: AsyncWrite + AsyncRead,
{
    pub fn new(prefix: Vec<u8>, stream: S) -> Self
    {
        let (read, write) = split(stream);
        log::trace!("Prefix: {:?}", prefix);
        Self {
            read: PrefixedRead {
                prefix: Some(prefix),
                read,
            },
            write,
        }
    }

    pub fn into_split(self) -> (PrefixedRead<S>, WriteHalf<S>)
    {
        (self.read, self.write)
    }
}

impl<S: AsyncWrite + AsyncRead + Unpin> PrefixedStream<S> {}

impl<S: AsyncRead> AsyncRead for PrefixedStream<S>
{
    fn poll_read(self: Pin<&mut Self>, cx: &mut Context, buf: &mut ReadBuf) -> Poll<Result<()>>
    {
        let inner_pin = unsafe { self.map_unchecked_mut(|s| &mut s.read) };
        inner_pin.poll_read(cx, buf)
    }
}

impl<S: AsyncRead> AsyncRead for PrefixedRead<S>
{
    fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context, buf: &mut ReadBuf)
        -> Poll<Result<()>>
    {
        if let Some(p) = &mut self.prefix {
            if p.is_empty() {
                // If the Vec was empty, we'll want to avoid returning zero bytes. Tokio considers
                // a return of zero bytes as stream having ended, which is not what is happening
                // here.
                self.prefix = None;
            } else if p.len() <= buf.remaining() {
                buf.put_slice(p);
                self.prefix = None;
                return Poll::Ready(Ok(()));
            } else {
                let mut taken = p.split_off(buf.remaining());
                std::mem::swap(&mut taken, p);
                buf.put_slice(&taken);
                return Poll::Ready(Ok(()));
            }
        }

        // `self` is pinned, which results in `s.stream` being pinned.
        // This makes the `map_unchecked_mut` safe here.
        let inner_pin = unsafe { self.map_unchecked_mut(|s| &mut s.read) };
        inner_pin.poll_read(cx, buf)
    }
}

impl<S: AsyncWrite> AsyncWrite for PrefixedStream<S>
{
    fn poll_write(self: Pin<&mut Self>, cx: &mut Context, buf: &[u8]) -> Poll<Result<usize>>
    {
        // `self` is pinned, which results in `s.stream` being pinned.
        // This makes the `map_unchecked_mut` safe here.
        let inner_pin = unsafe { self.map_unchecked_mut(|s| &mut s.write) };
        inner_pin.poll_write(cx, buf)
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<()>>
    {
        let inner_pin = unsafe { self.map_unchecked_mut(|s| &mut s.write) };
        inner_pin.poll_flush(cx)
    }

    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<()>>
    {
        let inner_pin = unsafe { self.map_unchecked_mut(|s| &mut s.write) };
        inner_pin.poll_shutdown(cx)
    }
}