message: s2n-tls: improve send buffering
diff: @@ -62,7 +62,7 @@ where
  }
  /// Removes all of the callback and context pointers from the connection
  - pub fn unset(self, connection: &mut Connection) -> Result<(), TransportError> {
  + pub fn unset(mut self, connection: &mut Connection) -> Result<(), TransportError> {
  unsafe {
  unsafe extern "C" fn secret_cb(
  _context: *mut c_void,
  @@ -101,6 +101,9 @@ where
  .set_receive_context(core::ptr::null_mut())
  .unwrap();
  + // Flush the send buffer before returning to the connection
  + self.flush();
  +
  if let Some(err) = self.err {
  return Err(err);
  }
  @@ -187,6 +190,9 @@ where
  }
  };
  + // Flush the send buffer before transitioning to the next phase
  + self.flush();
  +
  match self.state.tx_phase {
  HandshakePhase::Initial => {
  let keys = RingHandshakeCrypto::new(self.endpoint, aead_algo, pair)
  @@ -228,20 +234,44 @@ where
  /// Called when sending data
  fn on_write(&mut self, data: &[u8]) -> usize {
  - if self.send_buffer.capacity() < data.len() {
  - *self.send_buffer = BytesMut::with_capacity(SEND_BUFFER_CAPACITY);
  + // If this write would cause the current send buffer to reallocate,
  + // we should flush and create a new send buffer.
  + let remaining_capacity = self.send_buffer.capacity() - self.send_buffer.len();
  +
  + if remaining_capacity < data.len() {
  + // Flush the send buffer before reallocating it
  + self.flush();
  +
  + // ensure we only do one allocation for this write
  + let len = SEND_BUFFER_CAPACITY.max(data.len());
  +
  + debug_assert!(
  + self.send_buffer.is_empty(),
  + "dropping a send buffer with data will result in data loss"
  + );
  + *self.send_buffer = BytesMut::with_capacity(len);
  }
  +
  + // Write the current data to the send buffer
  + //
  + // NOTE: we don't immediately flush to the packet space since s2n-tls may do
  + // several small writes in a row.
  self.send_buffer.extend_from_slice(data);
  - let chunk = self.send_buffer.split_to(data.len());
  - let chunk = chunk.freeze();
  +
  + data.len()
  + }
  +
  + /// Flushes the send buffer into the current TX space
  + fn flush(&mut self) {
  + if !self.send_buffer.is_empty() {
  + let chunk = self.send_buffer.split().freeze();
  match self.state.tx_phase {
  HandshakePhase::Initial => self.context.send_initial(chunk),
  HandshakePhase::Handshake => self.context.send_handshake(chunk),
  HandshakePhase::Application => self.context.send_application(chunk),
  }
  -
  - data.len()
  + }
  }
  /// The function s2n-tls calls when it wants to receive data
  @@ -259,7 +289,7 @@ where
  }
  }
  - /// Called when sending data
  + /// Called when receiving data
  fn on_read(&mut self, data: &mut [u8]) -> usize {
  let max_len = Some(data.len());
  