| use crate::codec::Codec; |
| use crate::frame::{self, Reason, StreamId}; |
| |
| use bytes::Buf; |
| use std::io; |
| use std::task::{Context, Poll}; |
| use tokio::io::AsyncWrite; |
| |
| /// Manages our sending of GOAWAY frames. |
| #[derive(Debug)] |
| pub(super) struct GoAway { |
| /// Whether the connection should close now, or wait until idle. |
| close_now: bool, |
| /// Records if we've sent any GOAWAY before. |
| going_away: Option<GoingAway>, |
| /// Whether the user started the GOAWAY by calling `abrupt_shutdown`. |
| is_user_initiated: bool, |
| /// A GOAWAY frame that must be buffered in the Codec immediately. |
| pending: Option<frame::GoAway>, |
| } |
| |
| /// Keeps a memory of any GOAWAY frames we've sent before. |
| /// |
| /// This looks very similar to a `frame::GoAway`, but is a separate type. Why? |
| /// Mostly for documentation purposes. This type is to record status. If it |
| /// were a `frame::GoAway`, it might appear like we eventually wanted to |
| /// serialize it. We **only** want to be able to look up these fields at a |
| /// later time. |
| #[derive(Debug)] |
| pub(crate) struct GoingAway { |
| /// Stores the highest stream ID of a GOAWAY that has been sent. |
| /// |
| /// It's illegal to send a subsequent GOAWAY with a higher ID. |
| last_processed_id: StreamId, |
| |
| /// Records the error code of any GOAWAY frame sent. |
| reason: Reason, |
| } |
| |
| impl GoAway { |
| pub fn new() -> Self { |
| GoAway { |
| close_now: false, |
| going_away: None, |
| is_user_initiated: false, |
| pending: None, |
| } |
| } |
| |
| /// Enqueue a GOAWAY frame to be written. |
| /// |
| /// The connection is expected to continue to run until idle. |
| pub fn go_away(&mut self, f: frame::GoAway) { |
| if let Some(ref going_away) = self.going_away { |
| assert!( |
| f.last_stream_id() <= going_away.last_processed_id, |
| "GOAWAY stream IDs shouldn't be higher; \ |
| last_processed_id = {:?}, f.last_stream_id() = {:?}", |
| going_away.last_processed_id, |
| f.last_stream_id(), |
| ); |
| } |
| |
| self.going_away = Some(GoingAway { |
| last_processed_id: f.last_stream_id(), |
| reason: f.reason(), |
| }); |
| self.pending = Some(f); |
| } |
| |
| pub fn go_away_now(&mut self, f: frame::GoAway) { |
| self.close_now = true; |
| if let Some(ref going_away) = self.going_away { |
| // Prevent sending the same GOAWAY twice. |
| if going_away.last_processed_id == f.last_stream_id() && going_away.reason == f.reason() |
| { |
| return; |
| } |
| } |
| self.go_away(f); |
| } |
| |
| pub fn go_away_from_user(&mut self, f: frame::GoAway) { |
| self.is_user_initiated = true; |
| self.go_away_now(f); |
| } |
| |
| /// Return if a GOAWAY has ever been scheduled. |
| pub fn is_going_away(&self) -> bool { |
| self.going_away.is_some() |
| } |
| |
| pub fn is_user_initiated(&self) -> bool { |
| self.is_user_initiated |
| } |
| |
| /// Returns the going away info, if any. |
| pub fn going_away(&self) -> Option<&GoingAway> { |
| self.going_away.as_ref() |
| } |
| |
| /// Returns if the connection should close now, or wait until idle. |
| pub fn should_close_now(&self) -> bool { |
| self.pending.is_none() && self.close_now |
| } |
| |
| /// Returns if the connection should be closed when idle. |
| pub fn should_close_on_idle(&self) -> bool { |
| !self.close_now |
| && self |
| .going_away |
| .as_ref() |
| .map(|g| g.last_processed_id != StreamId::MAX) |
| .unwrap_or(false) |
| } |
| |
| /// Try to write a pending GOAWAY frame to the buffer. |
| /// |
| /// If a frame is written, the `Reason` of the GOAWAY is returned. |
| pub fn send_pending_go_away<T, B>( |
| &mut self, |
| cx: &mut Context, |
| dst: &mut Codec<T, B>, |
| ) -> Poll<Option<io::Result<Reason>>> |
| where |
| T: AsyncWrite + Unpin, |
| B: Buf, |
| { |
| if let Some(frame) = self.pending.take() { |
| if !dst.poll_ready(cx)?.is_ready() { |
| self.pending = Some(frame); |
| return Poll::Pending; |
| } |
| |
| let reason = frame.reason(); |
| dst.buffer(frame.into()).expect("invalid GOAWAY frame"); |
| |
| return Poll::Ready(Some(Ok(reason))); |
| } else if self.should_close_now() { |
| return match self.going_away().map(|going_away| going_away.reason) { |
| Some(reason) => Poll::Ready(Some(Ok(reason))), |
| None => Poll::Ready(None), |
| }; |
| } |
| |
| Poll::Ready(None) |
| } |
| } |
| |
| impl GoingAway { |
| pub(crate) fn reason(&self) -> Reason { |
| self.reason |
| } |
| } |