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
use crate::{handle::UnsafeHandle, session::Session};
use futures::{AsyncRead, AsyncWrite};
use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
use windows_sys::Win32::{
    Foundation::{FALSE, HANDLE, WAIT_ABANDONED_0, WAIT_EVENT, WAIT_OBJECT_0},
    System::Threading::{WaitForMultipleObjects, INFINITE},
};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum WaitingStopReason {
    Shutdown,
    Ready,
}

#[derive(Debug, Clone)]
enum ReadState {
    Waiting(Option<Arc<Mutex<blocking::Task<WaitingStopReason>>>>),
    Idle,
    Closed,
}

#[derive(Clone)]
pub struct AsyncSession {
    session: Arc<Session>,
    read_state: ReadState,
}

impl std::ops::Deref for AsyncSession {
    type Target = Session;

    fn deref(&self) -> &Self::Target {
        &self.session
    }
}

impl From<Arc<Session>> for AsyncSession {
    fn from(session: Arc<Session>) -> Self {
        Self {
            session,
            read_state: ReadState::Idle,
        }
    }
}

impl AsyncSession {
    fn wait_for_read(read_event: UnsafeHandle<HANDLE>, shutdown_event: UnsafeHandle<HANDLE>) -> WaitingStopReason {
        const WAIT_OBJECT_1: WAIT_EVENT = WAIT_OBJECT_0 + 1;
        const WAIT_ABANDONED_1: WAIT_EVENT = WAIT_ABANDONED_0 + 1;
        let handles = [shutdown_event.0, read_event.0];
        match unsafe { WaitForMultipleObjects(handles.len() as u32, &handles as _, FALSE, INFINITE) } {
            WAIT_OBJECT_0 | WAIT_ABANDONED_0 => WaitingStopReason::Shutdown,
            WAIT_OBJECT_1 => WaitingStopReason::Ready,
            WAIT_ABANDONED_1 => panic!("Read event deleted unexpectedly"),
            e => panic!("WaitForMultipleObjects returned unexpected value {:?}", e),
        }
    }
}

impl AsyncRead for AsyncSession {
    fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context, buf: &mut [u8]) -> Poll<std::io::Result<usize>> {
        loop {
            match &mut self.read_state {
                ReadState::Idle => match self.session.try_receive() {
                    Ok(Some(packet)) => {
                        let size = packet.bytes.len().min(buf.len());
                        buf[..size].copy_from_slice(&packet.bytes[..size]);
                        return Poll::Ready(Ok(size));
                    }
                    Ok(None) => {
                        let read_event = self.session.get_read_wait_event()?;
                        let shutdown_event = self.session.shutdown_event.get_handle();
                        let task = Arc::new(Mutex::new(blocking::unblock(move || {
                            Self::wait_for_read(read_event, shutdown_event)
                        })));
                        self.read_state = ReadState::Waiting(Some(task));
                    }
                    Err(err) => return Poll::Ready(Err(std::io::Error::new(std::io::ErrorKind::Other, err))),
                },
                ReadState::Waiting(task) => {
                    let task = match task.take() {
                        Some(task) => task,
                        None => return Poll::Pending,
                    };
                    let task_clone = task.clone();
                    let mut task_guard = match task_clone.lock() {
                        Ok(guard) => guard,
                        Err(e) => {
                            self.read_state = ReadState::Waiting(Some(task));
                            use std::io::{Error, ErrorKind};
                            return Poll::Ready(Err(Error::new(ErrorKind::Other, format!("Lock task failed: {}", e))));
                        }
                    };
                    self.read_state = match Pin::new(&mut *task_guard).poll(cx) {
                        Poll::Ready(WaitingStopReason::Shutdown) => ReadState::Closed,
                        Poll::Ready(WaitingStopReason::Ready) => ReadState::Idle,
                        Poll::Pending => ReadState::Waiting(Some(task)),
                    };
                    if let ReadState::Waiting(_) = self.read_state {
                        return Poll::Pending;
                    }
                }
                ReadState::Closed => return Poll::Ready(Ok(0)),
            }
        }
    }
}

impl AsyncWrite for AsyncSession {
    fn poll_write(self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &[u8]) -> Poll<std::io::Result<usize>> {
        let packet = self.session.allocate_send_packet(buf.len() as _)?;
        packet.bytes.copy_from_slice(buf);
        self.session.send_packet(packet);
        Poll::Ready(Ok(buf.len()))
    }

    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
        Poll::Ready(Ok(()))
    }

    fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
        self.session.shutdown()?;
        Poll::Ready(Ok(()))
    }
}