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

use std::io;

use tokio::io::unix::AsyncFd;

pub struct TokioSocket2 {
    io: AsyncFd<socket2::Socket>,
}

impl TokioSocket2 {
    pub fn new(socket: socket2::Socket) -> io::Result<Self> {
        let io = AsyncFd::new(socket)?;
        Ok(Self { io })
    }

    pub fn get_ref(&self) -> &socket2::Socket {
        self.io.get_ref()
    }

    pub async fn read<F: FnMut(&socket2::Socket) -> io::Result<R>, R>(
        &self,
        mut f: F
    ) -> io::Result<R> {
        loop {
            let mut guard = self.io.readable().await?;

            match guard.try_io(|io| f(io.get_ref())) {
                Ok(result) => {
                    return result;
                }
                Err(_would_block) => {
                    continue;
                }
            }
        }
    }

    pub async fn write<F: FnMut(&socket2::Socket) -> io::Result<R>, R>(
        &self,
        mut f: F
    ) -> io::Result<R> {
        loop {
            let mut guard = self.io.writable().await?;

            match guard.try_io(|io| f(io.get_ref())) {
                Ok(result) => {
                    return result;
                }
                Err(_would_block) => {
                    continue;
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::mem::MaybeUninit;
    use std::{ io, mem };
    use std::net::SocketAddr;

    use socket2::{ Domain, Protocol, Socket, Type };
    use tokio::{ net::TcpStream, io::{ AsyncReadExt, AsyncWriteExt } };

    use super::TokioSocket2;

    #[tokio::test]
    async fn test() -> io::Result<()> {
        let listen_socket = Socket::new(Domain::IPV4, Type::STREAM, Some(Protocol::TCP))?;
        listen_socket.set_nonblocking(true)?;
        listen_socket.set_reuse_address(true)?;
        listen_socket.bind(&SocketAddr::from(([127, 0, 0, 1], 0)).into())?;
        listen_socket.listen(1)?;

        let listener = TokioSocket2::new(listen_socket)?;

        eprintln!("Listener");

        let listen_addr = listener.get_ref().local_addr()?.as_socket().unwrap();

        let mut client = TcpStream::connect(listen_addr).await?;

        eprintln!("Client");

        let (server_socket, _) = listener.read(|socket| socket.accept()).await?;
        let server = TokioSocket2::new(server_socket)?;

        eprintln!("Server");

        client.write_all(b"ping").await?;

        eprintln!("Client write");

        let mut buf = [0; 4];
        let mut pos = 0;

        while pos < 4 {
            let n = server.read(|socket| {
                let buf = unsafe {
                    mem::transmute::<&mut [u8], &mut [MaybeUninit<u8>]>(&mut buf[pos..])
                };
                socket.recv(buf)
            }).await?;

            pos += n;
        }

        eprintln!("Server read");

        assert_eq!(&buf[..4], b"ping");

        server.write(|socket| socket.send(b"pong")).await?;

        eprintln!("Server write");

        client.read_exact(&mut buf[..4]).await?;

        eprintln!("Client read");

        assert_eq!(&buf[..4], b"pong");

        Ok(())
    }
}