fix bug that SSL response should only contain the type (even without length)

Author: Tianyi Chen

src/include/network/marshal.h

@@ -138,6 +138,8 @@ struct OutputPacket {
   size_t ptr;                   // ByteBuf cursor, which is used for get and put
   NetworkMessageType msg_type;  // header
 
+  bool single_type_pkt;         // there would be only a pkt type being written to the
+                                // buffer when this flag is true
   bool skip_header_write;  // whether we should write header to socket wbuf
   size_t write_ptr;        // cursor used to write packet content to socket wbuf
 
@@ -146,6 +148,7 @@ struct OutputPacket {
     buf.resize(BUFFER_INIT_SIZE);
     buf.shrink_to_fit();
     buf.clear();
+    single_type_pkt = false;
     len = ptr = write_ptr = 0;
     msg_type = NetworkMessageType::NULL_COMMAND;
     skip_header_write = true;

src/network/connection_handle.cpp

@@ -121,10 +121,10 @@ DEF_TRANSITION_GRAPH
   END_DEF
 
 END_DEF
-// clang-format on
+    // clang-format on
 
-void ConnectionHandle::StateMachine::Accept(Transition action,
-                                            ConnectionHandle &connection) {
+    void ConnectionHandle::StateMachine::Accept(Transition action,
+                                                ConnectionHandle &connection) {
   Transition next = action;
   while (next != Transition::NONE) {
     transition_result result = Delta_(current_state_, next);
@@ -466,16 +466,17 @@ WriteState ConnectionHandle::BufferWriteBytesHeader(OutputPacket *pkt) {
     wbuf_->buf[wbuf_->buf_ptr++] = type;
   }
 
-  // make len include its field size as well
-  len_nb = htonl(len + sizeof(int32_t));
+  if (!pkt->single_type_pkt) {
+    // make len include its field size as well
+    len_nb = htonl(len + sizeof(int32_t));
 
-  // TODO (Tianyi) Check why we need finish startup packet before
-  // append the bytes of this integer in network-byte order
-  std::copy(reinterpret_cast<uchar *>(&len_nb),
-            reinterpret_cast<uchar *>(&len_nb) + 4,
-            std::begin(wbuf_->buf) + wbuf_->buf_ptr);
-  // move the write buffer pointer and update size of the socket buffer
-  wbuf_->buf_ptr += sizeof(int32_t);
+    // append the bytes of this integer in network-byte order
+    std::copy(reinterpret_cast<uchar *>(&len_nb),
+              reinterpret_cast<uchar *>(&len_nb) + 4,
+              std::begin(wbuf_->buf) + wbuf_->buf_ptr);
+    // move the write buffer pointer and update size of the socket buffer
+    wbuf_->buf_ptr += sizeof(int32_t);
+  }
 
   wbuf_->buf_size = wbuf_->buf_ptr;
 
@@ -558,7 +559,7 @@ Transition ConnectionHandle::CloseSocket() {
     conn_SSL_context = nullptr;
   }
 
-  while(true) {
+  while (true) {
     int status = close(sock_fd_);
     if (status < 0) {
       // failed close
@@ -601,8 +602,7 @@ Transition ConnectionHandle::SSLHandshake() {
   // clear current thread's error queue before any OpenSSL call
   ERR_clear_error();
   int ssl_accept_ret = SSL_accept(conn_SSL_context);
-  if (ssl_accept_ret > 0)
-    return Transition::PROCEED;
+  if (ssl_accept_ret > 0) return Transition::PROCEED;
 
   int err = SSL_get_error(conn_SSL_context, ssl_accept_ret);
   int ecode = ERR_get_error();

src/network/postgres_protocol_handler.cpp

@@ -58,6 +58,7 @@ const std::unordered_map<std::string, std::string>
 
 PostgresProtocolHandler::PostgresProtocolHandler(tcop::TrafficCop *traffic_cop)
     : ProtocolHandler(traffic_cop),
+      init_stage_(true),
       txn_state_(NetworkTransactionStateType::IDLE) {}
 
 PostgresProtocolHandler::~PostgresProtocolHandler() {}
@@ -976,15 +977,16 @@ ProcessResult PostgresProtocolHandler::ProcessInitialPacket(InputPacket *pkt) {
   int32_t proto_version = PacketGetInt(pkt, sizeof(int32_t));
   LOG_INFO("protocol version: %d", proto_version);
 
+  force_flush_ = true;
   // TODO(Yuchen): consider more about return value
   if (proto_version == SSL_MESSAGE_VERNO) {
     LOG_TRACE("process SSL MESSAGE");
     std::unique_ptr<OutputPacket> response(new OutputPacket());
     bool ssl_able = (PelotonServer::GetSSLLevel() != SSLLevel::SSL_DISABLE);
     response->msg_type =
         ssl_able ? NetworkMessageType::SSL_YES : NetworkMessageType::SSL_NO;
+    response->single_type_pkt = true;
     responses_.push_back(std::move(response));
-    force_flush_ = true;
     return ssl_able ? ProcessResult::NEED_SSL_HANDSHAKE
                     : ProcessResult::COMPLETE;
   } else {