-
Notifications
You must be signed in to change notification settings - Fork 3
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Signed-off-by: iGxnon <igxnon@gmail.com>
- Loading branch information
Showing
8 changed files
with
93 additions
and
313 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,250 +1,126 @@ | ||
| use std::net::SocketAddr; | ||
| use std::pin::Pin; | ||
| use std::task::{ready, Context, Poll}; | ||
| use std::time::{SystemTime, UNIX_EPOCH}; | ||
|
|
||
| use bytes::Bytes; | ||
| use futures::{Future, Sink, SinkExt, Stream, StreamExt}; | ||
| use futures::Stream; | ||
| use log::debug; | ||
| use pin_project_lite::pin_project; | ||
|
|
||
| use crate::errors::{CodecError, Error}; | ||
| use crate::link::SharedLink; | ||
| use crate::packet::connected::FrameBody; | ||
| use crate::Message; | ||
| use crate::RoleContext; | ||
|
|
||
| pub(crate) trait HandleOnline: Sized { | ||
| fn handle_online<O>( | ||
| fn handle_online( | ||
| self, | ||
| write: O, | ||
| addr: SocketAddr, | ||
| client_guid: u64, | ||
| ) -> OnlineHandler<Self, O> | ||
| where | ||
| O: Sink<FrameBody, Error = CodecError>; | ||
| link: SharedLink, | ||
| ) -> OnlineHandler<Self>; | ||
| } | ||
|
|
||
| impl<F> HandleOnline for F | ||
| where | ||
| F: Stream<Item = FrameBody>, | ||
| { | ||
| fn handle_online<O>( | ||
| fn handle_online( | ||
| self, | ||
| write: O, | ||
| addr: SocketAddr, | ||
| client_guid: u64, | ||
| ) -> OnlineHandler<Self, O> | ||
| where | ||
| O: Sink<FrameBody, Error = CodecError>, | ||
| { | ||
| link: SharedLink, | ||
| ) -> OnlineHandler<Self> { | ||
| link.send_frame_body(FrameBody::ConnectionRequest { | ||
| client_guid, | ||
| request_timestamp: timestamp(), | ||
| use_encryption: false, | ||
| }); | ||
| OnlineHandler { | ||
| read: Some(self), | ||
| write: Some(write), | ||
| state: State::SendConnectionRequest(Some(FrameBody::ConnectionRequest { | ||
| client_guid, | ||
| request_timestamp: timestamp(), | ||
| use_encryption: false, | ||
| })), | ||
| frame: self, | ||
| state: State::WaitConnRes, | ||
| addr, | ||
| link, | ||
| role: RoleContext::Client { guid: client_guid }, | ||
| } | ||
| } | ||
| } | ||
|
|
||
| pin_project! { | ||
| pub(crate) struct OnlineHandler<I, O> { | ||
| read: Option<I>, | ||
| write: Option<O>, | ||
| pub(crate) struct OnlineHandler<F> { | ||
| #[pin] | ||
| frame: F, | ||
| state: State, | ||
| addr: SocketAddr, | ||
| link: SharedLink, | ||
| role: RoleContext, | ||
| } | ||
| } | ||
|
|
||
| enum State { | ||
| SendConnectionRequest(Option<FrameBody>), | ||
| WaitConnectionRequestReply(Option<FrameBody>), | ||
| SendNewIncomingConnection(FrameBody), | ||
| WaitConnRes, | ||
| Connected, | ||
| } | ||
|
|
||
| impl<I, O> Future for OnlineHandler<I, O> | ||
| impl<F> Stream for OnlineHandler<F> | ||
| where | ||
| I: Stream<Item = FrameBody> + Unpin, | ||
| O: Sink<FrameBody, Error = CodecError> + Sink<Message, Error = CodecError> + Unpin, | ||
| F: Stream<Item = FrameBody>, | ||
| { | ||
| type Output = Result<impl Stream<Item = Bytes> + Sink<Message, Error = Error>, Error>; | ||
| type Item = Bytes; | ||
|
|
||
| fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { | ||
| let this = self.project(); | ||
| let write = this.write.as_mut().unwrap(); | ||
| let read = this.read.as_mut().unwrap(); | ||
| fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { | ||
| let mut this = self.project(); | ||
| loop { | ||
| match this.state { | ||
| State::SendConnectionRequest(pack) => { | ||
| if let Err(err) = ready!(SinkExt::<FrameBody>::poll_ready_unpin(write, cx)) { | ||
| debug!("[client] SendConnectionRequest poll_ready error: {err}, retrying"); | ||
| continue; | ||
| } | ||
| if let Err(err) = write.start_send_unpin(pack.clone().unwrap()) { | ||
| debug!("[client] SendConnectionRequest start_send error: {err}, retrying"); | ||
| continue; | ||
| } | ||
| if let Err(err) = ready!(SinkExt::<FrameBody>::poll_flush_unpin(write, cx)) { | ||
| debug!("[client] SendConnectionRequest poll_flush error: {err}, retrying"); | ||
| continue; | ||
| } | ||
| *this.state = State::WaitConnectionRequestReply(pack.take()); | ||
| } | ||
| State::WaitConnectionRequestReply(pack) => { | ||
| let Some(body) = ready!(read.poll_next_unpin(cx)) else { | ||
| return Poll::Ready(Err(Error::ConnectionClosed)); | ||
| State::WaitConnRes => { | ||
| let Some(body) = ready!(this.frame.as_mut().poll_next(cx)) else { | ||
| return Poll::Ready(None); | ||
| }; | ||
| match body { | ||
| FrameBody::ConnectionRequestAccepted { | ||
| if let FrameBody::ConnectionRequestAccepted { | ||
| system_addresses, | ||
| accepted_timestamp, | ||
| .. | ||
| } = body | ||
| { | ||
| this.link.send_frame_body(FrameBody::NewIncomingConnection { | ||
| server_address: *this.addr, | ||
| system_addresses, | ||
| request_timestamp: timestamp(), | ||
| accepted_timestamp, | ||
| .. | ||
| } => { | ||
| *this.state = State::SendNewIncomingConnection( | ||
| FrameBody::NewIncomingConnection { | ||
| server_address: *this.addr, | ||
| system_addresses, | ||
| request_timestamp: timestamp(), | ||
| accepted_timestamp, | ||
| }, | ||
| ); | ||
| } | ||
| _ => { | ||
| debug!("[client] got unexpected packet {body:?} on WaitConnectionRequestReply, fallback to SendConnectionRequest"); | ||
| *this.state = State::SendConnectionRequest(pack.take()); | ||
| } | ||
| } | ||
| } | ||
| State::SendNewIncomingConnection(pack) => { | ||
| if let Err(err) = ready!(SinkExt::<FrameBody>::poll_ready_unpin(write, cx)) { | ||
| }); | ||
| *this.state = State::Connected; | ||
| debug!( | ||
| "[client] SendNewIncomingConnection poll_ready error: {err}, retrying" | ||
| "[{}] connected to server {addr:?}", | ||
| this.role, | ||
| addr = this.addr | ||
| ); | ||
| continue; | ||
| } | ||
| if let Err(err) = write.start_send_unpin(pack.clone()) { | ||
| debug!( | ||
| "[client] SendNewIncomingConnection start_send error: {err}, retrying" | ||
| ); | ||
| continue; | ||
| } | ||
| if let Err(err) = ready!(SinkExt::<FrameBody>::poll_flush_unpin(write, cx)) { | ||
| debug!( | ||
| "[client] SendNewIncomingConnection poll_flush error: {err}, retrying" | ||
| ); | ||
| continue; | ||
| } | ||
| return Poll::Ready(Ok(FilterIO { | ||
| read: this.read.take().unwrap(), | ||
| write: this.write.take().unwrap(), | ||
| state: FilterIOState::Serving, | ||
| })); | ||
| debug!("[{}] ignore packet {body:?} on WaitConnRes", this.role); | ||
| } | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
||
| pin_project! { | ||
| struct FilterIO<I, O> { | ||
| read: I, | ||
| write: O, | ||
| state: FilterIOState, | ||
| } | ||
| } | ||
|
|
||
| enum FilterIOState { | ||
| Serving, | ||
| SendPing, | ||
| Closed, | ||
| } | ||
|
|
||
| impl<I, O> Stream for FilterIO<I, O> | ||
| where | ||
| I: Stream<Item = FrameBody> + Unpin, | ||
| O: Sink<FrameBody, Error = CodecError> + Unpin, | ||
| { | ||
| type Item = Bytes; | ||
|
|
||
| fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { | ||
| let this = self.project(); | ||
| loop { | ||
| match this.state { | ||
| FilterIOState::Serving => { | ||
| let Some(body) = ready!(this.read.poll_next_unpin(cx)) else { | ||
| *this.state = FilterIOState::Closed; | ||
| continue; | ||
| State::Connected => { | ||
| let Some(body) = ready!(this.frame.as_mut().poll_next(cx)) else { | ||
| return Poll::Ready(None); | ||
| }; | ||
| match body { | ||
| FrameBody::DisconnectNotification => *this.state = FilterIOState::Closed, | ||
| FrameBody::DetectLostConnections => *this.state = FilterIOState::SendPing, | ||
| FrameBody::DetectLostConnections => { | ||
| this.link.send_frame_body(FrameBody::ConnectedPing { | ||
| client_timestamp: timestamp(), | ||
| }); | ||
| } | ||
| FrameBody::User(data) => return Poll::Ready(Some(data)), | ||
| _ => { | ||
| debug!("[client] ignore packet {body:?} on Connected",); | ||
| debug!("[{}] ignore packet {body:?} on Connected", this.role); | ||
| } | ||
| } | ||
| } | ||
| FilterIOState::SendPing => { | ||
| if let Err(err) = ready!(this.write.poll_ready_unpin(cx)) { | ||
| debug!("[client] SendPing poll_ready error: {err}, retrying"); | ||
| continue; | ||
| } | ||
| if let Err(err) = this.write.start_send_unpin(FrameBody::ConnectedPing { | ||
| client_timestamp: timestamp(), | ||
| }) { | ||
| debug!("[client] SendPing start_send error: {err}, retrying"); | ||
| continue; | ||
| } | ||
| if let Err(err) = ready!(this.write.poll_flush_unpin(cx)) { | ||
| debug!("[client] SendPing poll_flush error: {err}, retrying"); | ||
| continue; | ||
| } | ||
| *this.state = FilterIOState::Serving; | ||
| } | ||
| FilterIOState::Closed => return Poll::Ready(None), | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
||
| impl<I, O> Sink<Message> for FilterIO<I, O> | ||
| where | ||
| O: Sink<Message, Error = CodecError> + Unpin, | ||
| { | ||
| type Error = Error; | ||
|
|
||
| fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { | ||
| let this = self.project(); | ||
| if matches!(*this.state, FilterIOState::Closed) { | ||
| return Poll::Ready(Err(Error::ConnectionClosed)); | ||
| } | ||
| ready!(this.write.poll_ready_unpin(cx))?; | ||
| Poll::Ready(Ok(())) | ||
| } | ||
|
|
||
| fn start_send(self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> { | ||
| self.project().write.start_send_unpin(item)?; | ||
| Ok(()) | ||
| } | ||
|
|
||
| fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { | ||
| ready!(self.project().write.poll_flush_unpin(cx))?; | ||
| Poll::Ready(Ok(())) | ||
| } | ||
|
|
||
| fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { | ||
| let this = self.project(); | ||
| ready!(this.write.poll_close_unpin(cx))?; | ||
| *this.state = FilterIOState::Closed; | ||
| Poll::Ready(Ok(())) | ||
| } | ||
| } | ||
|
|
||
| fn timestamp() -> i64 { | ||
| SystemTime::now() | ||
| .duration_since(UNIX_EPOCH) | ||
| std::time::SystemTime::now() | ||
| .duration_since(std::time::UNIX_EPOCH) | ||
| .unwrap() | ||
| .as_millis() as i64 | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.