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
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, AsyncWrite, Error, ReadBuf};
use tokio::net::TcpStream;
#[cfg(unix)]
use tokio::net::UnixStream;

use std::io;

use crate::Addr;

/// A socket connected to an endpoint
#[derive(Debug)]
pub enum Stream {
    /// A TCP stream between a local and a remote socket.
    Inet(TcpStream),
    #[cfg(unix)]
    /// A connected Unix socket
    Unix(UnixStream),
}

impl From<TcpStream> for Stream {
    fn from(s: TcpStream) -> Stream {
        Stream::Inet(s)
    }
}

#[cfg(unix)]
impl From<UnixStream> for Stream {
    fn from(s: UnixStream) -> Stream {
        Stream::Unix(s)
    }
}

impl Stream {
    /// Opens a connection to a remote host.
    pub async fn connect(s: &Addr) -> io::Result<Stream> {
        match s {
            Addr::Inet(s) => TcpStream::connect(s).await.map(Stream::Inet),
            #[cfg(unix)]
            Addr::Unix(s) => UnixStream::connect(s).await.map(Stream::Unix),
        }
    }

    /// Returns the local address that this stream is bound to.
    pub fn local_addr(&self) -> io::Result<Addr> {
        match self {
            Stream::Inet(s) => s.local_addr().map(Addr::Inet),
            #[cfg(unix)]
            Stream::Unix(s) => s.local_addr().map(|e| e.into()),
        }
    }

    /// Returns the remote address that this stream is connected to.
    pub fn peer_addr(&self) -> io::Result<Addr> {
        match self {
            Stream::Inet(s) => s.peer_addr().map(Addr::Inet),
            #[cfg(unix)]
            Stream::Unix(s) => s.peer_addr().map(|e| e.into()),
        }
    }
}
impl AsyncRead for Stream {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<Result<(), Error>> {
        match &mut *self {
            Stream::Inet(s) => Pin::new(s).as_mut().poll_read(cx, buf),
            #[cfg(unix)]
            Stream::Unix(s) => Pin::new(s).as_mut().poll_read(cx, buf),
        }
    }
}
impl AsyncWrite for Stream {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context,
        buf: &[u8],
    ) -> Poll<Result<usize, Error>> {
        match &mut *self {
            Stream::Inet(s) => Pin::new(s).as_mut().poll_write(cx, buf),
            #[cfg(unix)]
            Stream::Unix(s) => Pin::new(s).as_mut().poll_write(cx, buf),
        }
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Error>> {
        match &mut *self {
            Stream::Inet(s) => Pin::new(s).as_mut().poll_flush(cx),
            #[cfg(unix)]
            Stream::Unix(s) => Pin::new(s).as_mut().poll_flush(cx),
        }
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Error>> {
        match &mut *self {
            Stream::Inet(s) => Pin::new(s).as_mut().poll_shutdown(cx),
            #[cfg(unix)]
            Stream::Unix(s) => Pin::new(s).as_mut().poll_shutdown(cx),
        }
    }
}



#[cfg(test)]
pub(crate) mod tests {
    use super::*;
    use std::net::SocketAddr;
    use tokio::io::{AsyncReadExt, AsyncWriteExt};
    use tokio::net::TcpListener;
    #[cfg(unix)]
    use tokio::net::UnixListener;
    use tokio::runtime::Builder;

    pub(crate) async fn local_socket_pair() -> Result<(TcpListener, Addr), std::io::Error> {
        let a: SocketAddr = "127.0.0.1:0".parse().unwrap();
        let app_listener = TcpListener::bind(a).await?;
        let a: Addr = app_listener.local_addr()?.into();
        Ok((app_listener, a))
    }

    #[test]
    fn tcp_connect() {
        let rt = Builder::new_current_thread().enable_all().build().unwrap();
        async fn mock_app(app_listener: TcpListener) {
            let (mut app_socket, _) = app_listener.accept().await.unwrap();
            let mut buf = [0u8; 32];
            let i = app_socket.read(&mut buf).await.unwrap();
            app_socket.write_all(&buf[..i]).await.unwrap();
        }

        async fn con() {
            let (app_listener, a) = local_socket_pair().await.unwrap();
            tokio::spawn(mock_app(app_listener));

            let mut s = Stream::connect(&a).await.expect("tcp connect failed");

            let data = b"1234";
            s.write_all(&data[..]).await.expect("tcp write failed");

            let mut buf = [0u8; 32];
            let i = s.read(&mut buf).await.expect("tcp read failed");
            assert_eq!(&buf[..i], &data[..]);
        }
        rt.block_on(con());
    }
    #[cfg(unix)]
    #[test]
    fn unix_connect() {
        use std::path::Path;

        let rt = Builder::new_current_thread().enable_all().build().unwrap();
        async fn mock_app(app_listener: UnixListener) {
            let (mut app_socket, _) = app_listener.accept().await.unwrap();
            let mut buf = [0u8; 32];
            let i = app_socket.read(&mut buf).await.unwrap();
            app_socket.write_all(&buf[..i]).await.unwrap();
        }

        async fn con() {
            let a: &Path = Path::new("/tmp/afcgi.sock");
            let app_listener = UnixListener::bind(a).unwrap();
            tokio::spawn(mock_app(app_listener));

            let a: Addr = "/tmp/afcgi.sock".parse().expect("unix parse failed");
            let mut s = Stream::connect(&a).await.expect("unix connect failed");

            let data = b"1234";
            s.write_all(&data[..]).await.expect("unix write failed");

            let mut buf = [0u8; 32];
            let i = s.read(&mut buf).await.expect("unix read failed");
            assert_eq!(&buf[..i], &data[..]);
        }
        rt.block_on(con());
        std::fs::remove_file("/tmp/afcgi.sock").unwrap();
    }
}