Minnow¶

template<TCPDatagramAdapter AdaptT>
void TCPMinnowSocket<AdaptT>::_initialize_TCP( const TCPConfig& config )
{
  _tcp.emplace( config );

  // Set up the event loop

  // There are three events to handle:
  //
  // 1) Incoming datagram received (needs to be given to TCPPeer::receive method)
  //
  // 2) Outbound bytes received from local application via a write()
  //    call (needs to be read from the local stream socket and
  //    given to TCPPeer)
  //
  // 3) Incoming bytes reassembled by the Reassembler
  //    (needs to be read from the inbound_stream and written
  //    to the local stream socket back to the application)

  // rule 1: read from filtered packet stream and dump into TCPConnection
  _eventloop.add_rule(
    "receive TCP segment from the network",
    _datagram_adapter.fd(),
    Direction::In,
    [&] {
      if ( auto seg = _datagram_adapter.read() ) {
        _tcp->receive( std::move( seg.value() ), [&]( auto x ) { _datagram_adapter.write( x ); } );
      }

      // debugging output:
      if ( _thread_data.eof() and _tcp.value().sender().sequence_numbers_in_flight() == 0 and not _fully_acked ) {
        std::cerr << "DEBUG: minnow outbound stream to " << _datagram_adapter.config().destination.to_string()
                  << " has been fully acknowledged.\n";
        _fully_acked = true;
      }
    },
    [&] { return _tcp->active(); } );

  // rule 2: read from pipe into outbound buffer
  _eventloop.add_rule(
    "push bytes to TCPPeer",
    _thread_data,
    Direction::In,
    [&] {
      std::string data;
      data.resize( _tcp->outbound_writer().available_capacity() );
      _thread_data.read( data );
      _tcp->outbound_writer().push( move( data ) );

      if ( _thread_data.eof() ) {
        _tcp->outbound_writer().close();
        _outbound_shutdown = true;

        // debugging output:
        std::cerr << "DEBUG: minnow outbound stream to " << _datagram_adapter.config().destination.to_string()
                  << " finished (" << _tcp.value().sender().sequence_numbers_in_flight() << " seqno"
                  << ( _tcp.value().sender().sequence_numbers_in_flight() == 1 ? "" : "s" )
                  << " still in flight).\n";
      }

      _tcp->push( [&]( auto x ) { _datagram_adapter.write( x ); } );
    },
    [&] {
      return ( _tcp->active() ) and ( not _outbound_shutdown )
             and ( _tcp->outbound_writer().available_capacity() > 0 );
    },
    [&] {
      _tcp->outbound_writer().close();
      _outbound_shutdown = true;
    },
    [&] {
      std::cerr << "DEBUG: minnow outbound stream had error.\n";
      _tcp->outbound_writer().set_error();
    } );

  // rule 3: read from inbound buffer into pipe
  _eventloop.add_rule(
    "read bytes from inbound stream",
    _thread_data,
    Direction::Out,
    [&] {
      Reader& inbound = _tcp->inbound_reader();
      // Write from the inbound_stream into
      // the pipe, handling the possibility of a partial
      // write (i.e., only pop what was actually written).
      if ( inbound.bytes_buffered() ) {
        const std::string_view buffer = inbound.peek();
        const auto bytes_written = _thread_data.write( buffer );
        inbound.pop( bytes_written );
      }

      if ( inbound.is_finished() or inbound.has_error() ) {
        _thread_data.shutdown( SHUT_WR );
        _inbound_shutdown = true;

        // debugging output:
        std::cerr << "DEBUG: minnow inbound stream from " << _datagram_adapter.config().destination.to_string()
                  << " finished " << ( inbound.has_error() ? "uncleanly.\n" : "cleanly.\n" );
      }
    },
    [&] {
      return _tcp->inbound_reader().bytes_buffered()
             or ( ( _tcp->inbound_reader().is_finished() or _tcp->inbound_reader().has_error() )
                  and not _inbound_shutdown );
    },
    [&] {},
    [&] {
      std::cerr << "DEBUG: minnow inbound stream had error.\n";
      _tcp->inbound_reader().set_error();
    } );
}

template<TCPDatagramAdapter AdaptT>
void TCPMinnowSocket<AdaptT>::_tcp_main()
{
  try {
    if ( not _tcp.has_value() ) {
      throw std::runtime_error( "no TCP" );
    }
    _tcp_loop( [] { return true; } );
    shutdown( SHUT_RDWR );
    if ( not _tcp.value().active() ) {
      std::cerr << "DEBUG: minnow TCP connection finished "
                << ( _tcp->inbound_reader().has_error() ? "uncleanly.\n" : "cleanly.\n" );
    }
    _tcp.reset();
  } catch ( const std::exception& e ) {
    std::cerr << "Exception in TCPConnection runner thread: " << e.what() << "\n";
    throw e;
  }
}

//! \param[in] condition is a function returning true if loop should continue
template<TCPDatagramAdapter AdaptT>
void TCPMinnowSocket<AdaptT>::_tcp_loop( const std::function<bool()>& condition )
{
  auto base_time = timestamp_ms();
  while ( condition() ) {
    auto ret = _eventloop.wait_next_event( TCP_TICK_MS );
    if ( ret == EventLoop::Result::Exit or _abort ) {
      break;
    }

    if ( not _tcp.has_value() ) {
      throw std::runtime_error( "_tcp_loop entered before TCPPeer initialized" );
    }

    if ( _tcp.value().active() ) {
      const auto next_time = timestamp_ms();
      _tcp.value().tick( next_time - base_time, [&]( auto x ) { _datagram_adapter.write( x ); } );
      _datagram_adapter.tick( next_time - base_time );
      base_time = next_time;
    }
  }
}

class TCPPeer
{
  auto make_send( const auto& transmit )
  {
    return [&]( const TCPSenderMessage& x ) { send( x, transmit ); };
  }

public:
  void receive( TCPMessage msg, const TransmitFunction& transmit )
  {
    if ( not active() ) {
      return;
    }

    // Record time in case this peer has to linger after streams finish.
    time_of_last_receipt_ = cumulative_time_;

    // If SenderMessage occupies a sequence number, make sure to reply.
    need_send_ |= ( msg.sender.sequence_length() > 0 );

    // If SenderMessage is a "keep-alive" (with intentionally invalid seqno), make sure to reply.
    // (N.B. orthodox TCP rules require a reply on any unacceptable segment.)
    const auto our_ackno = receiver_.send().ackno;
    need_send_ |= ( our_ackno.has_value() and msg.sender.seqno + 1 == our_ackno.value() );

    // Did the inbound stream finish before the outbound stream? If so, no need to linger after streams finish.
    if ( receiver_.writer().is_closed() and not sender_.reader().is_finished() ) {
      linger_after_streams_finish_ = false;
    }

    // Give incoming TCPSenderMessage to receiver.
    receiver_.receive( std::move( msg.sender ) );

    // Give incoming TCPReceiverMessage to sender.
    sender_.receive( msg.receiver );

    // Send reply if needed.
    if ( need_send_ ) {
      send( sender_.make_empty_message(), transmit );
    }
  }


private:
  void send( const TCPSenderMessage& sender_message, const TransmitFunction& transmit )
  {
    TCPMessage msg { sender_message, receiver_.send() };
    transmit( std::move( msg ) );
    need_send_ = false;
  }
};

void bidirectional_stream_copy( Socket& socket, string_view peer_name )
{
  constexpr size_t buffer_size = 1048576;

  EventLoop _eventloop {};
  FileDescriptor _input { STDIN_FILENO };
  FileDescriptor _output { STDOUT_FILENO };
  ByteStream _outbound { buffer_size };
  ByteStream _inbound { buffer_size };
  bool _outbound_shutdown { false };
  bool _inbound_shutdown { false };

  socket.set_blocking( false );
  _input.set_blocking( false );
  _output.set_blocking( false );

  // rule 1: read from stdin into outbound byte stream
  _eventloop.add_rule(
    "read from stdin into outbound byte stream",
    _input,
    Direction::In,
    [&] {
      string data;
      data.resize( _outbound.writer().available_capacity() );
      _input.read( data );
      _outbound.writer().push( move( data ) );
      if ( _input.eof() ) {
        _outbound.writer().close();
      }
    },
    [&] {
      return !_outbound.has_error() and !_inbound.has_error() and ( _outbound.writer().available_capacity() > 0 )
             and !_outbound.writer().is_closed();
    },
    [&] { _outbound.writer().close(); },
    [&] {
      cerr << "DEBUG: Outbound stream had error from source.\n";
      _outbound.set_error();
      _inbound.set_error();
    } );

  // rule 2: read from outbound byte stream into socket
  _eventloop.add_rule(
    "read from outbound byte stream into socket",
    socket,
    Direction::Out,
    [&] {
      if ( _outbound.reader().bytes_buffered() ) {
        _outbound.reader().pop( socket.write( _outbound.reader().peek() ) );
      }
      if ( _outbound.reader().is_finished() ) {
        socket.shutdown( SHUT_WR );
        _outbound_shutdown = true;
        cerr << "DEBUG: Outbound stream to " << peer_name << " finished.\n";
      }
    },
    [&] {
      return _outbound.reader().bytes_buffered() or ( _outbound.reader().is_finished() and not _outbound_shutdown );
    },
    [&] { _outbound.writer().close(); },
    [&] {
      cerr << "DEBUG: Outbound stream had error from destination.\n";
      _outbound.set_error();
      _inbound.set_error();
    } );

  // rule 3: read from socket into inbound byte stream
  _eventloop.add_rule(
    "read from socket into inbound byte stream",
    socket,
    Direction::In,
    [&] {
      string data;
      data.resize( _inbound.writer().available_capacity() );
      socket.read( data );
      _inbound.writer().push( move( data ) );
      if ( socket.eof() ) {
        _inbound.writer().close();
      }
    },
    [&] {
      return !_inbound.has_error() and !_outbound.has_error() and ( _inbound.writer().available_capacity() > 0 )
             and !_inbound.writer().is_closed();
    },
    [&] { _inbound.writer().close(); },
    [&] {
      cerr << "DEBUG: Inbound stream had error from source.\n";
      _outbound.set_error();
      _inbound.set_error();
    } );

  // rule 4: read from inbound byte stream into stdout
  _eventloop.add_rule(
    "read from inbound byte stream into stdout",
    _output,
    Direction::Out,
    [&] {
      if ( _inbound.reader().bytes_buffered() ) {
        _inbound.reader().pop( _output.write( _inbound.reader().peek() ) );
      }
      if ( _inbound.reader().is_finished() ) {
        _output.close();
        _inbound_shutdown = true;
        cerr << "DEBUG: Inbound stream from " << peer_name << " finished"
             << ( _inbound.has_error() ? " uncleanly.\n" : ".\n" );
      }
    },
    [&] {
      return _inbound.reader().bytes_buffered() or ( _inbound.reader().is_finished() and not _inbound_shutdown );
    },
    [&] { _inbound.writer().close(); },
    [&] {
      cerr << "DEBUG: Inbound stream had error from destination.\n";
      _outbound.set_error();
      _inbound.set_error();
    } );

  // loop until completion
  while ( true ) {
    if ( EventLoop::Result::Exit == _eventloop.wait_next_event( -1 ) ) {
      return;
    }
  }
}

template<class T>
concept TCPDatagramAdapter = requires( T a, TCPMessage seg ) {
  {
    a.write( seg )
  } -> std::same_as<void>;

  {
    a.read()
  } -> std::same_as<std::optional<TCPMessage>>;
};

//! \brief A FD adapter for IPv4 datagrams read from and written to a TUN device
class TCPOverIPv4OverTunFdAdapter : public TCPOverIPv4Adapter
{
private:
  TunFD _tun;

public:
  //! Construct from a TunFD
  explicit TCPOverIPv4OverTunFdAdapter( TunFD&& tun ) : _tun( std::move( tun ) ) {}

  //! Attempts to read and parse an IPv4 datagram containing a TCP segment related to the current connection
  std::optional<TCPMessage> read();

  //! Creates an IPv4 datagram from a TCP segment and writes it to the TUN device
  void write( const TCPMessage& seg ) { _tun.write( serialize( wrap_tcp_in_ip( seg ) ) ); }

  //! Access the underlying TUN device
  explicit operator TunFD&() { return _tun; }

  //! Access the underlying TUN device
  explicit operator const TunFD&() const { return _tun; }

  //! Access underlying file descriptor
  FileDescriptor& fd() { return _tun; }
};

optional<TCPMessage> TCPOverIPv4OverTunFdAdapter::read()
{
  vector<string> strs( 2 );
  strs.front().resize( IPv4Header::LENGTH );
  _tun.read( strs );

  InternetDatagram ip_dgram;
  const vector<string> buffers = { strs.at( 0 ), strs.at( 1 ) };
  if ( parse( ip_dgram, buffers ) ) {
    return unwrap_tcp_in_ip( ip_dgram );
  }
  return {};
}

class TCPPeer
{

public:
  explicit TCPPeer( const TCPConfig& cfg ) : cfg_( cfg ) {}

  Writer& outbound_writer() { return sender_.writer(); }
  Reader& inbound_reader() { return receiver_.reader(); }

  /* Passthrough methods */
  void push( const TransmitFunction& transmit ) { sender_.push( make_send( transmit ) ); }

  void receive( TCPMessage msg, const TransmitFunction& transmit );

  // Testing interface
  const TCPReceiver& receiver() const { return receiver_; }
  const TCPSender& sender() const { return sender_; }

private:
  TCPConfig cfg_;
  TCPSender sender_ { ByteStream { cfg_.send_capacity }, cfg_.isn, cfg_.rt_timeout };
  TCPReceiver receiver_ { Reassembler { ByteStream { cfg_.recv_capacity } } };

  void send( const TCPSenderMessage& sender_message, const TransmitFunction& transmit );
};

void Writer::push( string data )
{
  // Your code here.
  size_t datalen = data.size();
  size_t write_avail = Writer::available_capacity();

  if ( Writer::is_closed() || write_avail == 0 || data.empty() ) {
    return;
  }
  if ( datalen > write_avail ) {
    data.resize( write_avail );
  }
  pushed_ += data.size();
  buffered_ += data.size();
  data_.emplace( std::move( data ) );
}

string_view Reader::peek() const
{
  // Your code here.
  return data_.empty() ? string_view {} : string_view { data_.front() }.substr( read_index_ );
}

void Reader::pop( uint64_t len )
{
  // Your code here.
  buffered_ -= len;
  poped_ += len;
  while ( len != 0U ) {
    uint64_t size = data_.front().size() - read_index_;
    if ( len < size ) {
      read_index_ += len;
      break;
    }
    // 需要向前移动一个string
    data_.pop();
    read_index_ = 0;
    len -= size;
  }
}

void Reassembler::insert( uint64_t first_index, string data, bool is_last_substring )
{
  // Your code here.
  const auto try_close = [&]() noexcept -> void {
    if ( end_idx_.has_value() && end_idx_.value() == next_unassem_idx_ ) {
      output_.writer().close();
    }
  };

  if ( data.empty() ) {
    if ( !end_idx_.has_value() && is_last_substring ) {
      end_idx_.emplace( first_index );
    }
    return try_close();
  }

  if ( writer().is_closed() || writer().available_capacity() == 0 ) {
    return;
  }

  uint64_t unassembled_index { next_unassem_idx_ };
  uint64_t unacceptable_index { unassembled_index + writer().available_capacity() };
  // 将不再[unassembled_index, unacceptable_index)的字符串进行处理
  if ( first_index + size( data ) <= unassembled_index || first_index >= unacceptable_index ) {
    return; // Out of ranger
  }
  if ( first_index + size( data ) > unacceptable_index ) { // Remove unacceptable bytes
    data.resize( unacceptable_index - first_index );
    is_last_substring = false;
  }
  if ( first_index < unassembled_index ) { // Remove poped/buffered bytes
    data.erase( 0, unassembled_index - first_index );
    first_index = unassembled_index;
  }

  if ( !end_idx_.has_value() and is_last_substring ) {
    end_idx_.emplace( first_index + size( data ) );
  }

  // 对字符串进行分割
  // TODO: 可以对字符串移除进行优化
  auto upper { split( first_index + size( data ) ) };
  auto lower { split( first_index ) };
  // 对[lower, upper)的迭代器进行视图转换，然后更新reassemble的大小
  ranges::for_each( ranges::subrange( lower, upper ) | views::values,
                    [&]( const auto& str ) { reassemble_len_ -= str.size(); } );
  reassemble_len_ += size( data );
  buffer_.emplace_hint( buffer_.erase( lower, upper ), first_index, move( data ) );

  // 向writer中写入已经排好序的报文
  while ( !buffer_.empty() ) {
    auto&& [index, packet] { *buffer_.begin() };
    if ( index != next_unassem_idx_ ) {
      break;
    }

    reassemble_len_ -= size( packet );
    next_unassem_idx_ += size( packet );
    output_.writer().push( move( packet ) );
    buffer_.erase( buffer_.begin() );
  }
  return try_close();
}