1 use crate::codec::Codec; 2 use crate::frame::{self, Reason, StreamId}; 3 4 use bytes::Buf; 5 use std::io; 6 use std::task::{Context, Poll}; 7 use tokio::io::AsyncWrite; 8 9 /// Manages our sending of GOAWAY frames. 10 #[derive(Debug)] 11 pub(super) struct GoAway { 12 /// Whether the connection should close now, or wait until idle. 13 close_now: bool, 14 /// Records if we've sent any GOAWAY before. 15 going_away: Option<GoingAway>, 16 /// Whether the user started the GOAWAY by calling `abrupt_shutdown`. 17 is_user_initiated: bool, 18 /// A GOAWAY frame that must be buffered in the Codec immediately. 19 pending: Option<frame::GoAway>, 20 } 21 22 /// Keeps a memory of any GOAWAY frames we've sent before. 23 /// 24 /// This looks very similar to a `frame::GoAway`, but is a separate type. Why? 25 /// Mostly for documentation purposes. This type is to record status. If it 26 /// were a `frame::GoAway`, it might appear like we eventually wanted to 27 /// serialize it. We **only** want to be able to look up these fields at a 28 /// later time. 29 /// 30 /// (Technically, `frame::GoAway` should gain an opaque_debug_data field as 31 /// well, and we wouldn't want to save that here to accidentally dump in logs, 32 /// or waste struct space.) 33 #[derive(Debug)] 34 pub(crate) struct GoingAway { 35 /// Stores the highest stream ID of a GOAWAY that has been sent. 36 /// 37 /// It's illegal to send a subsequent GOAWAY with a higher ID. 38 last_processed_id: StreamId, 39 40 /// Records the error code of any GOAWAY frame sent. 41 reason: Reason, 42 } 43 44 impl GoAway { 45 pub fn new() -> Self { 46 GoAway { 47 close_now: false, 48 going_away: None, 49 is_user_initiated: false, 50 pending: None, 51 } 52 } 53 54 /// Enqueue a GOAWAY frame to be written. 55 /// 56 /// The connection is expected to continue to run until idle. 57 pub fn go_away(&mut self, f: frame::GoAway) { 58 if let Some(ref going_away) = self.going_away { 59 assert!( 60 f.last_stream_id() <= going_away.last_processed_id, 61 "GOAWAY stream IDs shouldn't be higher; \ 62 last_processed_id = {:?}, f.last_stream_id() = {:?}", 63 going_away.last_processed_id, 64 f.last_stream_id(), 65 ); 66 } 67 68 self.going_away = Some(GoingAway { 69 last_processed_id: f.last_stream_id(), 70 reason: f.reason(), 71 }); 72 self.pending = Some(f); 73 } 74 75 pub fn go_away_now(&mut self, f: frame::GoAway) { 76 self.close_now = true; 77 if let Some(ref going_away) = self.going_away { 78 // Prevent sending the same GOAWAY twice. 79 if going_away.last_processed_id == f.last_stream_id() && going_away.reason == f.reason() 80 { 81 return; 82 } 83 } 84 self.go_away(f); 85 } 86 87 pub fn go_away_from_user(&mut self, f: frame::GoAway) { 88 self.is_user_initiated = true; 89 self.go_away_now(f); 90 } 91 92 /// Return if a GOAWAY has ever been scheduled. 93 pub fn is_going_away(&self) -> bool { 94 self.going_away.is_some() 95 } 96 97 pub fn is_user_initiated(&self) -> bool { 98 self.is_user_initiated 99 } 100 101 /// Returns the going away info, if any. 102 pub fn going_away(&self) -> Option<&GoingAway> { 103 self.going_away.as_ref() new(codec: Codec<T, Prioritized<B>>, config: Config) -> Connection<T, P, B>104 } 105 106 /// Returns if the connection should close now, or wait until idle. 107 pub fn should_close_now(&self) -> bool { 108 self.pending.is_none() && self.close_now 109 } 110 111 /// Returns if the connection should be closed when idle. 112 pub fn should_close_on_idle(&self) -> bool { 113 !self.close_now 114 && self 115 .going_away 116 .as_ref() 117 .map(|g| g.last_processed_id != StreamId::MAX) 118 .unwrap_or(false) 119 } 120 121 /// Try to write a pending GOAWAY frame to the buffer. 122 /// 123 /// If a frame is written, the `Reason` of the GOAWAY is returned. 124 pub fn send_pending_go_away<T, B>( 125 &mut self, 126 cx: &mut Context, 127 dst: &mut Codec<T, B>, 128 ) -> Poll<Option<io::Result<Reason>>> 129 where 130 T: AsyncWrite + Unpin, 131 B: Buf, 132 { 133 if let Some(frame) = self.pending.take() { 134 if !dst.poll_ready(cx)?.is_ready() { 135 self.pending = Some(frame); 136 return Poll::Pending; 137 } 138 139 let reason = frame.reason(); 140 dst.buffer(frame.into()).expect("invalid GOAWAY frame"); 141 142 return Poll::Ready(Some(Ok(reason))); 143 } else if self.should_close_now() { 144 return match self.going_away().map(|going_away| going_away.reason) { set_target_window_size(&mut self, size: WindowSize)145 Some(reason) => Poll::Ready(Some(Ok(reason))), 146 None => Poll::Ready(None), 147 }; 148 } 149 set_initial_window_size(&mut self, size: WindowSize) -> Result<(), UserError>150 Poll::Ready(None) 151 } 152 } 153 154 impl GoingAway { 155 pub(crate) fn reason(&self) -> Reason { 156 self.reason 157 } 158 } 159